GoでABC205 Dを２つのやり方で解く

概要

ABC205 D問題をGoで二つのやり方で解きます。一つは公式解法で、もう一つは「答えで二分探索」による解法です。ちなみに、サイトに掲載されている解法には、クエリをソートして二分探索なしで解くやり方がありますが、TLE(Time Limit Exceeded) エラーが解消できなかったため、断念しました。

\def\middlemid{\;\middle|\;}
\def\set#1{\left\{{#1}\right\}}
\def\setin#1#2{\left\{{#1} \middlemid {#2}\right\}}
\def\sub#1#2{{#1}_{#2}}

公式の想定解法

ざっくりした方針

$\sub{A}{i}$ごとにその$\sub{A}{i}$より小さく、かつ$A$に含まれない正の整数の個数を$\sub{C}{i}$に記録します。例えば、$A=\set{1,5,6,7}$とすると、$\sub{C}{2}$は、$\sub{A}{2}=6$より小さく、かつ$1,5$を含まないため、$3$となります。
少し飛躍がありますが、答えの$A$に含まれない正の整数を小さい方から数えて$K$番目の数は, その$K$に近い$\sub{C}{i}$を探して、うまく補正すると求めることができます。

コード

package main

import (
	"bufio"
	"fmt"
	"os"
)

var reader = bufio.NewReader(os.Stdin)
var writer = bufio.NewWriter(os.Stdout)

func abs(x int) int {
  if x > 0 {
    return x
  } else {
    return x * -1
  }
}

func isOK(a *[]int, index, key int) bool {
	return (*a)[index] >= key
}

func binarySearch(a *[]int, key int) int {
	ng := -1
	ok := len(*a)

	for abs(ok-ng) > 1 {
		mid := (ok + ng) / 2

		if isOK(a, mid, key) {
			ok = mid
		} else {
			ng = mid
		}
	}
	return ok
}

func main() {
	defer writer.Flush()

	var n, q int
	fmt.Fscan(reader, &n, &q)

	var a = make([]int, n)
	for i := 0; i < n; i++ {
		fmt.Fscan(reader, &a[i])
	}

	var c = make([]int, n)
	for i := 0; i < n; i++ {
		if i == 0 {
			c[i] = a[i] - 1
			continue
		}
		c[i] = c[i-1] + a[i] - a[i-1] - 1
	}

	var k = make([]int, q)
	for i := 0; i < q; i++ {
		fmt.Fscan(reader, &k[i])
	}

	for i := 0; i < q; i++ {
		ans := 0
		if c[n-1] < k[i] {
			ans = k[i] - c[n-1] + a[n-1]
		} else {
			idx := binarySearch(&c, k[i])
			ans = a[idx] - 1 - (c[idx] - k[i])
		}
		fmt.Fprintf(writer, "%d\n", ans)
	}
}

「答えで二分探索」による解法

ざっくりした方針

ある正の整数 $m$ から $A$に含まれる$m$以下の整数の個数を引いた数は、$m$以下の$A$に含まれない正の整数の個数になります。これが「答え」なので、あとは二分探索で「答え」が$k$以上を満たす最小の$m$を求めることができます。

コード

package main
 
import (
	"bufio"
	"fmt"
	"math"
	"os"
)
 
var reader = bufio.NewReader(os.Stdin)
var writer = bufio.NewWriter(os.Stdout)
 
const inf = math.MaxInt64
 
func abs(x int) int {
	if x < 0 {
		return x * -1
	} else {
		return x
	}
}
 
func isOK(a *[]int, index, key int) bool {
	return (*a)[index] > key
}
 
func binarySearch(a *[]int, key int) int {
	ng := -1
	ok := len(*a)
 
	for abs(ok-ng) > 1 {
		mid := (ok + ng) / 2
 
		if isOK(a, mid, key) {
			ok = mid
		} else {
			ng = mid
		}
	}
	return ok
}
 
func main() {
	defer writer.Flush()
 
	var n, q int
	fmt.Fscan(reader, &n, &q)
 
	var a = make([]int, n)
	for i := 0; i < n; i++ {
		fmt.Fscan(reader, &a[i])
	}
 
	for i := 0; i < q; i++ {
		var k int
		fmt.Fscan(reader, &k)
		ng := 0
		ok := inf
 
		for abs(ok-ng) > 1 {
			mid := (ok + ng) / 2
			idx := binarySearch(&a, mid)
 
			if mid-idx >= k {
				ok = mid
			} else {
				ng = mid
			}
		}
 
		fmt.Fprintf(writer, "%d\n", ok)
	}
 
}