概要

本棚から「現在左からT番目の本」の耐久値を減らし、耐久値が0になったら取り除く操作を効率的に処理する問題です。本が取り除かれると詰められるため、「現在のk番目」が元のどの本かを高速に求める必要があります。

考察

素朴なアプローチの問題点

配列で本棚を管理し、本を取り除くたびに要素を削除する方法が最も直感的です。しかし、配列の途中から要素を削除するには \(O(N)\) かかり、操作が \(Q\) 回あるため全体で \(O(NQ)\) となります。\(N, Q\) が最大 \(2 \times 10^5\) なので、\(4 \times 10^{10}\) 程度の計算量となり TLE になります。

重要な気づき

本が取り除かれて詰まるという操作は、「現在残っている本の中で左からk番目のものを見つける」 操作に帰着できます。これは 「k番目の要素を高速に検索する」 という典型的な問題で、Binary Indexed Tree（BIT / Fenwick Tree） を用いた k-th element クエリ で解けます。

具体的には、各本の元の位置 \(i\) に対して「まだ残っていれば 1、取り除かれていれば 0」という値を BIT で管理します。すると、BIT 上の累積和が \(k\) 以上になる最小の位置を求めることで、「現在 \(k\) 番目の本の元の位置」が \(O(\log N)\) で分かります。

アルゴリズム

1. 初期化: 長さ \(N\) の BIT を作り、すべての位置に 1 をセットする。残り冊数 remaining = N。
2. 各クエリ \(T_j\) に対して:
   • \(T_j > \text{remaining}\) なら何もしない。
   • そうでなければ、BIT 上で \(T_j\) 番目の 1 がある位置（＝元の本の位置 \(p\)）を二分探索（binary lifting）で求める。
   • \(D[p]\) を 1 減らす。
   • \(D[p] = 0\) になったら、BIT の位置 \(p\) を 0 に更新し、remaining を 1 減らす。
   • remaining を出力する。

BIT 上の k-th element の求め方（binary lifting）

BIT の構造を利用して、累積和が \(k\) 以上になる最小の位置を \(O(\log N)\) で求めます。2のべき乗を大きい方から試し、累積和が \(k\) 未満なら進む、という貪欲法です。

具体例: \(N=5\), \(D = [2, 1, 3, 1, 2]\) で \(T=3\) のとき、BIT 上で3番目の1がある位置は元の位置3です。もし位置2の本が既に取り除かれていれば、BIT上の3番目は元の位置4になります。

計算量

• 時間計算量: \(O((N + Q) \log N)\)
  • BIT の初期化に \(O(N \log N)\)
  • 各クエリで k-th element の検索に \(O(\log N)\)、BIT の更新に \(O(\log N)\)
• 空間計算量: \(O(N)\)
  • BIT 配列と耐久値配列

実装のポイント

• BIT は 1-indexed で管理するのが定石です。元の本の位置も 1-indexed で扱い、耐久値配列 \(D\) は 0-indexed なので D[pos - 1] でアクセスします。

• find_kth 関数では、BIT のサイズ以上の最小の2のべき乗から開始して binary lifting を行います。これは通常の二分探索より定数倍が軽く高速です。

• 入力は stdin.buffer.read() で一括読み込みし、出力も配列にためて最後に一括出力することで I/O のオーバーヘッドを削減しています。

  ソースコード

import sys
from sys import stdin

def main():
    input = stdin.buffer.read().split()
    idx = 0
    N = int(input[idx]); idx += 1
    Q = int(input[idx]); idx += 1
    
    D = [0] * N
    for i in range(N):
        D[i] = int(input[idx]); idx += 1
    
    # BIT (Fenwick Tree) to track which positions still have books
    # We store 1 for each position that still has a book
    # To find the k-th book, we do a binary search on the BIT
    
    size = N + 1
    bit = [0] * size
    
    def update(i, delta):
        # 1-indexed
        while i < size:
            bit[i] += delta
            i += i & (-i)
    
    def query(i):
        # prefix sum [1..i]
        s = 0
        while i > 0:
            s += bit[i]
            i -= i & (-i)
        return s
    
    def find_kth(k):
        # Find the smallest position p such that prefix_sum(p) >= k
        # Using binary lifting on BIT
        pos = 0
        r = 1
        while r < size:
            r <<= 1
        r >>= 1
        cur_sum = 0
        while r > 0:
            if pos + r < size and cur_sum + bit[pos + r] < k:
                pos += r
                cur_sum += bit[pos]
            r >>= 1
        return pos + 1
    
    # Initialize BIT: all N books present
    for i in range(1, N + 1):
        update(i, 1)
    
    remaining = N
    
    out = []
    for _ in range(Q):
        T = int(input[idx]); idx += 1
        if T > remaining:
            out.append(str(remaining))
        else:
            # Find the T-th remaining book's original position
            pos = find_kth(T)  # 1-indexed original position
            D[pos - 1] -= 1
            if D[pos - 1] == 0:
                update(pos, -1)
                remaining -= 1
            out.append(str(remaining))
    
    sys.stdout.write('\n'.join(out) + '\n')

main()

この解説は claude4.6opus-thinking によって生成されました。