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)
    }

}