ABC440 DをPythonで解いた

どうもこんにちは！

ABC440のDが公式の解説を理解するのに時間がかかってしまい、Pythonでのコーディングも結構な時間を使ってようやくという感じでした。。
せっかくなので追加で整理した結果をまとめておきます。
Cより前の問題はこちらの記事をご確認ください。

問題と公式の解説のリンク

問題は以下のリンクから。

D - Forbidden List 2 -

問題

整数がn個のリストAとx,yが与えられ、x以上の整数でリストに含まれない値のうち、小さい方からy番目の値を出力するという問題。
例えばリストが1 2 3 9 16、x=12, y=4とします。x(12)以上でリストに含まれない値は12,13,14,15,17,18・・・となり、y(4)番目の値を出力なので15が答えとなります。
なお、テストケースとしてはx,yの組がq個与えられます。

考え方とコード

リスト内の要素ごとに、条件を満たす数がいくつあるかに着目して考えます。リスト内でx以上で最小の値のインデックスをsとすると、インデックスがs以上のリストの値は条件を満たす数があるので、これがy個以上ある値で最小のインデックスtを見つければ計算できるとなります。
このsとtはそれぞれ2分探索で見つけることができます。

まとめると、以下のアルゴリズムを実装することになります。

1. 与えられるリストをソート
2. 要素の値の最大値をリストの末尾に入れる
3. リスト内でx以上かつ最小の値を2分探索で取得
4. $A_t-x+1 - (t-s+1) >= y$ となる最小のtを2分探索で取得
5. $x+(y-1)+t-s$を出力

import bisect

n,q = map(int,input().split())

# 2分探索で条件を満たす境界が必ずあるようにするために、ソートしたリストに要素の上限を追加
a = [int(x) for x in input().split()]
a.sort()
a += [10**9+1]

for _ in range(q):
    x,y = map(int,input().split())

    # リスト内でx以上で最小である要素のインデックスを取得
    s = bisect.bisect_left(a,x)

    # At-x+1 - (t-s+1) >= y となる最小のtを2分探索で取得
    # t=sとなる可能性もあるため、ngはs-1から開始
    # 最終的に、ngは上記の計算結果がy未満、okはy以上となる要素のインデックスが入る
    ng, ok = s-1, n
    while ng + 1 < ok:
        mid = (ng + ok) // 2
        if a[mid]- x + 1 - (mid - s + 1) >= y:
            ok = mid
        else:
            ng = mid
    t = ok

    print(x + y-1 + t-s)

ではでは。