ABC205 D - Kth Excluded

https://atcoder.jp/contests/abc205/tasks/abc205_d (opens in a new tab)

茶色上位。二分探索。

正の整数列 A について、各々の要素以下の数で、かつそれ以前の整数列に出現しない整数の個数をカウントしておく。

a = list(map(int,input().split()))
a_lower = [e-i for i,e in enumerate(a,1)]

a_lowerは単調増加する。

例えば入力が1 3 5 8 9 10 12 15 17 30だとすると、a_lowerは0 1 2 4 4 4 5 7 8 20となる。 以下はこの例をもとに考える。

K = 3のときを考える。 a_lower上で3が2と4の間に位置することから、それに対応するaでの5と8の間に解は存在する。a_lowerの2に対応する、aでの5から3-2=1分進んだ6が解となる。

K = 4のときを考える。 a_lower上で4が2と4の間に位置することから、それに対応するaでの5と8にの間に解は存在する。a_lowerの2に対応する、aでの5から4-2=2分進んだ7が解となる。

K = 7のときを考える。 a_lower上で7が5と7の間に位置することから、それに対応するaでの12と15にの間に解は存在する。a_lowerの5に対応する、aでの12から7-5=2分進んだ14が解となる。

以上よりターゲット値で二分探索(左)を行い、得られた値をiとすると、インデックスi-1に対応するaの値から、ターゲット値からインデックスi-1に対応するa_lowerに値を引いたものを加えればよいこととなる。

from bisect import bisect_left
k = <target>
i = bisect_left(a_lower,k)
print(a[i-1] + (k - a_lower[i-1]))

尚、インデックスを-1 することから、ターゲットがa_lower[0]以下の値であった場合、即ち二分探索して、0 が返ってくるような場合はは別途処理が必要である。

from bisect import bisect_left
n,q = map(int,input().split())
a = list(map(int,input().split()))
a_lower = [e-i-1 for i,e in enumerate(a)]
for _ in range(q):
    k = int(input())
    if k <= a_lower[0]:
        print(k)
    else:
        i = bisect_left(a_lower,k)
        print(a[i-1] + (k-a_lower[i-1]))