AtCoder Beginner Contest 205 D - Kth Excluded【Python】

AtCoder

https://atcoder.jp/contests/abc205/tasks/abc205_d

AtCoder ProblemsのRecommendationで Difficulty: 713、Solve Probability:35%でした。
A1​,A2​,…,AN​のいずれとも異なる正整数が 1,2,4,8,9,10,11,…の場合は、 [1, 4, 8]のように連続するものについてはその最初の値のみ管理、いくつ連続するかについて[2, 1, 10**18]のように管理し、後者の累積和についても求め、クエリに対し二分探索で答えを求めていきました。

import bisect
N, Q = map(int, input().split())
A = list(map(int, input().split()))
# A1<200b>,A2<200b>,…,AN<200b>のいずれとも異なる正整数のうち連続するものの最初の値
start_int = []
# いくつ連続するか
count = []
# いくつ連続するかの累積和
cum_count = [0]
pre_a = 0
for i in range(N):
    if A[i] - pre_a > 1:
        start_int.append(pre_a + 1)
        cur_count = A[i] - pre_a - 1
        count.append(cur_count)
    if i == N - 1:
        start_int.append(A[i] + 1)
        cur_count = 10 ** 18
        count.append(cur_count)
    pre_a = A[i]
for i in range(len(count)):
    cum_count.append(cum_count[i] + count[i])
for i in range(Q):
    k = int(input())
    idx = bisect.bisect_left(cum_count, k)
    ans = start_int[idx-1] + k - cum_count[idx-1] - 1
    print(ans)