AtCoder-ABC205 D - Kth Excluded【Python解答例】

AtCoder Beginner Contest205のD問題についてPythonの解答例を記事にしていきます。
AtCoder Beginner Contest 205 - AtCoder

AtCoder Beginner Contest205 D - Kth Excluded

D - Kth Excluded

問題文

長さ N の正整数列 A=(A1,A2,…,AN) と Q 個のクエリが与えられます。

i(1≤i≤Q) 番目のクエリでは、正整数 Ki が与えられるので、A1,A2,…,AN のいずれとも異なる正整数のうち、小さい方から数えて Ki 番目のものを求めてください。

制約

・1≤N,Q≤10^5
・1≤A1

解答例

import bisect

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

A = list(map(int, input().split()))

K = []
for _ in range(q):
    k = int(input())
    K.append(k)

C = [0] * n
for i in range(n):
    C[i] = A[i] - i - 1

for i in range(q):
    if C[n-1] < K[i]:
        print(A[n-1] + (K[i] - C[n-1]))
    else:
        tmp = bisect.bisect_left(C,K[i])
        print(A[tmp] - C[tmp] + K[i] - 1)

解説

与えられる正整数列Aと異なる正の整数でK番目に小さい数を答える問題です。

C[ i ]：Aのi番目の要素以下の数において、Aに含まれない正整数の個数を格納した配列Cを用意しました。

公式解説にあるように2つの場合分けを行います。

C[ n ] < K[ i ] のとき、A[ n - 1 ]以上になりますので、A[n-1] + (K[i] - C[n-1])が解答になります。
(K[i] - C[n-1])がAnより大きい正の整数の個数になりますのでA[ n - 1 ] に足した値が答えです。

C[ n ] >= K[ i ] のとき、C[ i ] >= Kとなる最小のiを二分探索で求めて、tmpとしました。
A[tmp] - C[tmp] + K[i] - 1が求める解答になります。

上記のように場合分けしてクエリ毎に出力すればOKでした。

ABC205の関連記事はこちら
ebisuke33.hatenablog.com
ebisuke33.hatenablog.com