ABC205 D – Kth Excluded｜かえるのほんだな

はじめに

競プロの精進のために学びがあった問題についてまとめています。
目標は AtCoder 水色！

今のレートは茶色です。

問題

ABC205 D – Kth Excludedを解きました。

難易度は茶色。

考察

NとQの制約が10^5なので、各クエリの検査をO(N)で行ってしまうと間に合わない可能性が高い。

また、Aiの範囲が1<= Ai <= 10^18なので、各Aに対してのO(N)の検査も間に合わない。

そのため、各クエリに対する整数列Aに含まれない数の検査は、2分探索を用いる。

具体的には、「A以外の数列において、左からK番目の値がXであることが成り立つ」場合のXに対して、答えの範囲で2分探索を行う。

また、「A以外の数列において、左からK番目の値がXであることが成り立つ」かどうかの判定のためには、左からK番目までに何個のAが含まれるかを2分探索でカウントする必要がある。

ここで、「左からK番目の値がXであることが成り立つ」のは「A以外の数列」においてであるので、xがAの中に存在しないことを確認しなくてはいけない。

コード

from bisect import bisect_left, bisect_right
N, Q = map(int,input().split())
A = list(map(int,input().split()))
A.sort()
MAX = pow(10, 18) + N + 10

def search_x(A, k):
    # x 以下の範囲に A の要素が x-k 個あるかを考える
    left = 1
    right = MAX
    while right > left:
        mid = (right + left) // 2

        # mid までに A が何個存在しているか
        r = bisect_right(A, mid, lo=0, hi=len(A))
        if mid - r == k:
            if r > 0:
                if A[r-1] != mid:
                    return mid
            else:
                return mid

        if mid - r >= k:
            right = mid

        else:
            left = mid

for i in range(Q):
    k = int(input())
    print(search_x(A, k))