公式

E - 本棚の整理 / Organizing the Bookshelf 解説 by admin

Gemini 3.1 Pro (Thinking)

概要

本棚から条件を満たす本を取り除き、空いた隙間を詰める操作を繰り返す問題です。各操作の後に残っている本の冊数を求めます。

考察

まず、問題文の通りに配列(リスト)を使って素朴にシミュレーションする方法を考えてみましょう。 配列から要素を削除して左に詰める操作は、削除した位置より右にあるすべての要素を移動させる必要があるため、1回の操作につき最大で \(O(N)\) の時間がかかります。これを \(Q\) 回繰り返すと全体の計算量は \(O(QN)\) となり、今回の制約(\(N, Q \le 2 \times 10^5\))では実行時間制限(TLE)に引っかかってしまいます。

この問題を高速に解くための重要な気づきは、「本を実際に移動させるのではなく、各初期位置に本が『存在する(1)』か『存在しない(0)』かを管理する」というアプローチです。

初期状態ではすべての位置に本があるため、[1, 1, 1, ..., 1] という状態です。もし左から2番目の本が取り除かれたら、[1, 0, 1, ..., 1] となります。 このとき、「現在本棚にある左から \(k\) 番目の本」は、「値の累積和がちょうど \(k\) になるような最小のインデックス」として特定することができます。

このように、 1. 1点更新(本を取り除く = 1を0にする) 2. 累積和の取得・検索(左から \(k\) 番目の位置を見つける)

の2つの操作を高速に行えるデータ構造を用いれば、効率的に問題を解くことができます。これには Binary Indexed Tree (BIT / Fenwick Tree)セグメント木 が適しています。

アルゴリズム

ここでは、実装が軽くて定数倍も高速な Binary Indexed Tree (BIT) を用いた解法を説明します。

  1. 初期化:

    • 長さ \(N\) の BIT を用意し、すべての位置 \(1, 2, \dots, N\)1 を設定します。
    • 本の耐久値の配列 \(D\) を用意します。
    • 現在の冊数 \(M\)\(N\) で初期化します。
  2. クエリの処理: 各操作 \(T_j\) について以下の手順を行います。

    • もし \(T_j > M\) なら、指定された位置に本は存在しないため、何もせずに \(M\) を出力して次の操作へ進みます。
    • \(T_j \le M\) の場合、BIT 上で二分探索を行い、「累積和が \(T_j\) となる位置(インデックス)」を求めます。これが、現在左から \(T_j\) 番目にある本の初期状態での位置です。
    • 見つけた位置にある本の耐久値 \(D\)\(1\) 減らします。
    • もし耐久値が \(0\) になった場合、その本は取り除かれるため、BIT のその位置の値を \(1\) 減らして 0 に更新します。同時に、現在の冊数 \(M\)\(1\) 減らします。
    • 操作後の冊数 \(M\) を出力します。

計算量

  • 時間計算量: \(O(N + Q \log N)\)
    • BITの初期化は \(O(N)\) で可能です。
    • 各クエリにおいて、BIT上の二分探索と値の更新にかかる時間は \(O(\log N)\) です。これが \(Q\) 回行われるため、クエリ処理全体で \(O(Q \log N)\) となります。
  • 空間計算量: \(O(N)\)
    • 本の耐久値を管理する配列と、BIT用の配列を持つため \(O(N)\) のメモリを使用します。

実装のポイント

  • BIT上の二分探索: 通常の二分探索とBITの sum 関数を組み合わせると \(O(\log^2 N)\) かかってしまいますが、ダブリング(2のべき乗の幅で上位ビットから決定していく手法)を用いることで、BITを直接探索し \(O(\log N)\) で目的のインデックスを見つけることができます。

  • 高速な入出力: Pythonでは入力サイズが大きい場合、input() をループで回すと遅くなります。sys.stdin.read().split() を使って入力を一括で読み込むことで、I/Oによるボトルネックを解消できます。

    ソースコード

import sys

def solve():
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    N = int(input_data[0])
    Q = int(input_data[1])
    
    D = [int(x) for x in input_data[2:N+2]]
    T = [int(x) for x in input_data[N+2:]]
    
    tree = [0] * (N + 1)
    for i in range(1, N + 1):
        tree[i] += 1
        nxt = i + (i & -i)
        if nxt <= N:
            tree[nxt] += tree[i]
            
    max_pow = 1
    while max_pow * 2 <= N:
        max_pow *= 2
        
    M = N
    out = []
    
    for t in T:
        if t <= M:
            idx = 0
            k = t
            step = max_pow
            while step > 0:
                next_idx = idx + step
                if next_idx <= N and tree[next_idx] < k:
                    idx = next_idx
                    k -= tree[idx]
                step //= 2
            
            idx += 1
            D[idx - 1] -= 1
            if D[idx - 1] == 0:
                M -= 1
                i = idx
                while i <= N:
                    tree[i] -= 1
                    i += i & -i
        
        out.append(str(M))
        
    print('\n'.join(out))

if __name__ == '__main__':
    solve()

この解説は gemini-3.1-pro-thinking によって生成されました。

投稿日時:
最終更新: