公式

E - 図書館の蔵書検索 / Library Book Search 解説 by admin

Claude 4.6 Opus (Thinking)

概要

棚番号が指定区間内にあり、かつページ数が指定値以上である本の冊数を高速に求める問題です。各クエリの答えから \(K\) を引いた値(負なら \(0\))を出力します。

考察

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

各クエリに対して、棚番号 \(L_j\) から \(R_j\) までのすべての本を走査し、ページ数が \(T_j\) 以上かどうかを1冊ずつ確認する方法が最も単純です。しかし、最悪の場合1回のクエリで \(M\) 冊すべてを調べることになり、\(Q\) 回のクエリで合計 \(O(MQ)\) かかります。\(M, Q\) がともに \(10^5\) のとき \(10^{10}\) となり、制限時間に間に合いません。

重要な気づき

この問題は2次元の条件でフィルタリングしています:

  1. 棚番号が区間 \([L, R]\) に含まれる
  2. ページ数\(T\) 以上

棚番号の区間クエリはセグメント木が得意とする形です。さらに、各ノードに対してページ数のリストをソート済みで持っておけば、「\(T\) 以上の本が何冊あるか」を二分探索で高速に求められます。この組み合わせが Merge Sort Tree(マージソートツリー) です。

アルゴリズム

Merge Sort Tree とは

セグメント木の各ノードに、そのノードが担当する区間に属する値のソート済みリストを格納するデータ構造です。

構築

  1. セグメント木のサイズを \(N\) 以上の最小の2べきとする。
  2. 各葉ノードは1つの棚に対応し、その棚に置かれている本のページ数をソート済みリストとして持つ。
  3. 内部ノードは、2つの子ノードのリストをマージして作る(マージソートと同じ要領)。

例えば、棚1に [100, 300]、棚2に [200] があれば、これらを管理する親ノードには [100, 200, 300] が格納されます。

クエリ処理

棚区間 \([L, R]\) でページ数 \(T\) 以上の本の冊数を求めるには:

  1. セグメント木上で区間 \([L, R]\)\(O(\log N)\) 個のノードに分解する。
  2. 各ノードのソート済みリストに対して、bisect_left\(T\) 以上の要素数を \(O(\log M)\) で求める。
  3. それらを合計して \(C_j\) とし、\(\max(C_j - K, 0)\) を出力する。

具体例

棚が4つ、本が次のように配置されているとします: - 棚1: ページ数 [50, 200] - 棚2: ページ数 [300] - 棚3: ページ数 [100, 400] - 棚4: ページ数 [150]

クエリ \(L=1, R=3, T=150\) の場合、棚1〜3の本のうちページ数150以上は [200, 300, 400] の3冊です。\(K=1\) なら出力は \(\max(3-1, 0) = 2\) です。

計算量

  • 構築の時間計算量: \(O(M \log N)\) — 各要素は木の \(O(\log N)\) 個のノードに含まれ、マージは全体で \(O(M \log N)\)
  • クエリの時間計算量: 1回あたり \(O(\log^2 N)\)\(O(\log N)\) 個のノードそれぞれで二分探索 \(O(\log M)\)
  • 全体の時間計算量: \(O(M \log N + Q \log^2 N)\)
  • 空間計算量: \(O(M \log N)\) — 各要素が \(O(\log N)\) 個のノードに格納される

実装のポイント

  • セグメント木のサイズは \(N\) 以上の2べきにする。棚番号(1-indexed)をそのまま葉のインデックスに変換(size + s - 1)する。

  • 内部ノードの構築は、Pythonの extend よりも手動マージのほうが定数倍が良い。

  • クエリでは半開区間 \([lo, hi)\) に変換して、lo & 1hi & 1 で端のノードを拾う標準的なセグメント木の走査を行う。

  • bisect_left(arr, T) はソート済みリストで \(T\) 以上の最初の位置を返すので、len(arr) - bisect_left(arr, T)\(T\) 以上の要素数になる。

  • 入力は sys.stdin.buffer.read() で一括読み込みし、出力も '\n'.join(...) でまとめて書き出すことで高速化する。

    ソースコード

import sys
from bisect import bisect_left

def main():
    input_data = sys.stdin.buffer.read().split()
    idx = 0
    N = int(input_data[idx]); idx += 1
    M = int(input_data[idx]); idx += 1
    Q = int(input_data[idx]); idx += 1
    K = int(input_data[idx]); idx += 1
    
    # Merge sort tree: for each node in a segment tree, store sorted list of page counts
    # Segment tree size
    size = 1
    while size < N:
        size *= 2
    
    # Each leaf corresponds to a shelf (1-indexed shelf -> 0-indexed leaf)
    tree = [[] for _ in range(2 * size)]
    
    for i in range(M):
        s = int(input_data[idx]); idx += 1
        d = int(input_data[idx]); idx += 1
        tree[size + s - 1].append(d)
    
    # Sort leaves
    for i in range(size, 2 * size):
        tree[i].sort()
    
    # Build internal nodes by merging children
    for i in range(size - 1, 0, -1):
        # Merge two sorted lists
        left = tree[2 * i]
        right = tree[2 * i + 1]
        merged = []
        li, ri = 0, 0
        ll, rl = len(left), len(right)
        while li < ll and ri < rl:
            if left[li] <= right[ri]:
                merged.append(left[li])
                li += 1
            else:
                merged.append(right[ri])
                ri += 1
        if li < ll:
            merged.extend(left[li:])
        if ri < rl:
            merged.extend(right[ri:])
        tree[i] = merged
    
    # Query: count of elements >= T in shelf range [L, R] (1-indexed)
    # In segment tree, shelf s is at index (size + s - 1)
    # We query range [L-1, R-1] in 0-indexed leaves, which is [size+L-1, size+R-1] in tree
    
    out = []
    for j in range(Q):
        L = int(input_data[idx]); idx += 1
        R = int(input_data[idx]); idx += 1
        T = int(input_data[idx]); idx += 1
        
        # Count elements >= T in range [L, R]
        count = 0
        lo = size + L - 1  # left boundary (inclusive)
        hi = size + R - 1  # right boundary (inclusive)
        lo_node = lo
        hi_node = hi
        
        # Convert to half-open: [lo, hi+1)
        lo = lo
        hi = hi + 1
        
        while lo < hi:
            if lo & 1:
                # lo is a right child, include it
                arr = tree[lo]
                count += len(arr) - bisect_left(arr, T)
                lo += 1
            if hi & 1:
                hi -= 1
                arr = tree[hi]
                count += len(arr) - bisect_left(arr, T)
            lo >>= 1
            hi >>= 1
        
        result = max(count - K, 0)
        out.append(str(result))
    
    sys.stdout.write('\n'.join(out) + '\n')

main()

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

投稿日時:
最終更新: