Official

E - 本棚の整理 / Organizing the Bookshelf Editorial by admin Editorial for printing

gemini-3.5-flash-thinking

概要

本棚の特定の区間に本を追加し(区間加算)、特定の棚の本をすべて取り除き(一点クリア)、特定の区間の本の合計冊数を求める(区間和取得)という、3種類のクエリを高速に処理する問題です。

考察

素朴なアプローチ

各棚の冊数を通常の配列で管理することを考えます。 - 作業 1(区間加算): 配列の \(l\) 番目から \(r\) 番目までに \(1\) を加える。最悪で \(O(N)\)。 - 作業 2(一点クリア): \(p\) 番目の要素を \(0\) にする。 \(O(1)\)。 - 作業 3(区間和取得): 配列の \(l\) 番目から \(r\) 番目までの合計をループで求める。最悪で \(O(N)\)。

この方法では、1回あたりの作業に最大で \(O(N)\) の時間がかかります。作業の回数 \(Q\) も最大で \(2 \times 10^5\) であるため、全体の計算量は \(O(NQ)\) となり、実行時間制限に間に合いません(TLEとなります)。

高速化のアイデア

区間加算と区間和取得を高速に行うための代表的なデータ構造として、遅延評価セグメント木(Lazy Segment Tree)があります。 遅延評価セグメント木を使うと、区間加算と区間和取得をそれぞれ \(O(\log N)\) で行うことができます。

しかし、本問には「作業 2:特定の位置 \(p\) の値を \(0\) にする」という操作があります。通常の遅延評価セグメント木は区間加算のみをサポートしているため、この「一点クリア」をどう処理するかがポイントになります。

一点クリアの解決策

セグメント木において、ある一点 \(p\) の値を書き換えるためには、その位置に影響を与えている保留中の加算操作(遅延値)をすべて解消(伝播)させる必要があります。

具体的には、以下の手順を踏むことで一点クリアを \(O(\log N)\) で実現できます。 1. 遅延値の伝播(Top-down): 根(ルート)から葉(位置 \(p\))に向かって、ルートから \(p\) へのパス上にあるノードの遅延値を子ノードへ順に押し下げ(push)ていきます。これにより、葉 \(p\) の値が最新の正しい値になります。 2. 値の書き換え: 葉 \(p\) の値を \(0\) に変更します。 3. 区間和の再計算(Bottom-up): 変更された葉 \(p\) から根に向かって、親ノードの値を「左の子の値 + 右の子の値」で更新していきます。

これにより、遅延評価セグメント木の整合性を保ったまま、一点クリアを正しく行うことができます。


アルゴリズム

遅延評価セグメント木を用いて、各クエリを以下のように処理します。

1. セグメント木の構築

各ノードについて、以下の情報を管理します。 - tree_val[i]: ノード \(i\) が表す区間の現在の合計冊数。 - tree_len[i]: ノード \(i\) がカバーする区間の長さ(棚の数)。区間加算の際に、加算する値にこの長さを掛けて区間和に足し合わせるために使用します。 - lazy[i]: ノード \(i\) に保留されている遅延加算値。

2. 各クエリの処理

  • 作業 1(区間加算): 区間 \([l, r)\) に対して、通常通りの遅延評価セグメント木の区間加算を行います。必要なノードに遅延値 lazy を加算し、区間和 tree_val を更新します。
  • 作業 2(一点クリア):
    1. 根から葉 \(p\) へのパス上にあるノードの lazy を、上から順に push して伝播させます。
    2. 葉 \(p\) に対応するノードの tree_val を \(0\) にします。
    3. 葉 \(p\) から根に向かって、親ノードの値を「左の子 + 右の子」で再計算します。
  • 作業 3(区間和取得): 区間 \([l, r)\) に対して、通常通りの遅延評価セグメント木の区間和取得を行います。必要に応じて遅延値を伝播させながら、該当する区間の和を足し合わせて返します。

計算量

  • 時間計算量:

    • セグメント木の初期化: \(O(N)\)
    • 各クエリ(作業 1, 2, 3)の処理: 全て \(O(\log N)\)
    • 全体での時間計算量: \(O(N + Q \log N)\) これは \(N, Q \le 2 \times 10^5\) の制約下で、実行時間制限(通常 2.0 秒)に十分間に合います。
  • 空間計算量:

    • セグメント木のサイズは \(N\) 以上の最小の \(2\) のべき乗を \(M\) としたとき、サイズ \(2M\) の配列がいくつか必要になります。
    • 全体での空間計算量: \(O(N)\) メモリ制限にも余裕で収まります。

実装のポイント

  • 1-based index と 0-based index の変換: 問題文では棚の番号が \(1\) から \(N\) まで(1-based)で与えられますが、プログラム内部では \(0\) から \(N-1\) まで(0-based)で処理するために、入力の \(l, p\) から \(1\) を引いています(区間の右端 \(r\) は半開区間 \([l, r)\) として扱うため、そのまま用います)。

  • ビット演算による高速化: セグメント木の親ノードへの移動(i >>= 1)や、子ノードへの移動(2 * i や 2 * i + 1)をビット演算で行うことで、Pythonでも高速に動作するように実装されています。

  • tree_len の活用: 区間に \(x\) を加算するとき、その区間の和は「(区間の長さ) \(\times x\)」だけ増加します。各ノードがカバーする区間の長さを tree_len 配列に予め計算して持たせておくことで、push 時や区間加算時の計算が \(O(1)\) で行えます。

    ソースコード

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])
    
    size = 1
    log = 0
    while size < N:
        size <<= 1
        log += 1
        
    tree_val = [0] * (2 * size)
    tree_len = [0] * (2 * size)
    lazy = [0] * size
    
    for i in range(N):
        tree_len[size + i] = 1
    for i in range(size - 1, 0, -1):
        tree_len[i] = tree_len[2 * i] + tree_len[2 * i + 1]
        
    def push(k):
        f = lazy[k]
        if f != 0:
            lk = 2 * k
            tree_val[lk] += f * tree_len[lk]
            if lk < size:
                lazy[lk] += f
            rk = lk + 1
            tree_val[rk] += f * tree_len[rk]
            if rk < size:
                lazy[rk] += f
            lazy[k] = 0

    def point_clear(p):
        g = p + size
        for i in range(log, 0, -1):
            push(g >> i)
        tree_val[g] = 0
        g >>= 1
        while g > 0:
            tree_val[g] = tree_val[2 * g] + tree_val[2 * g + 1]
            g >>= 1

    def range_add(l, r, x):
        if l >= r:
            return
        l += size
        r += size
        
        for i in range(log, 0, -1):
            mask = (1 << i) - 1
            if l & mask:
                push(l >> i)
            if r & mask:
                push((r - 1) >> i)
                
        l_opt, r_opt = l, r
        while l < r:
            if l & 1:
                tree_val[l] += x * tree_len[l]
                if l < size:
                    lazy[l] += x
                l += 1
            if r & 1:
                r -= 1
                tree_val[r] += x * tree_len[r]
                if r < size:
                    lazy[r] += x
            l >>= 1
            r >>= 1
            
        l, r = l_opt, r_opt
        for i in range(1, log + 1):
            mask = (1 << i) - 1
            if l & mask:
                p = l >> i
                tree_val[p] = tree_val[2 * p] + tree_val[2 * p + 1]
            if r & mask:
                p = (r - 1) >> i
                tree_val[p] = tree_val[2 * p] + tree_val[2 * p + 1]

    def range_sum(l, r):
        if l >= r:
            return 0
        l += size
        r += size
        
        for i in range(log, 0, -1):
            mask = (1 << i) - 1
            if l & mask:
                push(l >> i)
            if r & mask:
                push((r - 1) >> i)
                
        res = 0
        while l < r:
            if l & 1:
                res += tree_val[l]
                l += 1
            if r & 1:
                r -= 1
                res += tree_val[r]
            l >>= 1
            r >>= 1
        return res

    out = []
    idx = 2
    len_data = len(input_data)
    while idx < len_data:
        t = input_data[idx]
        if t == '1':
            l = int(input_data[idx+1]) - 1
            r = int(input_data[idx+2])
            range_add(l, r, 1)
            idx += 3
        elif t == '2':
            p = int(input_data[idx+1]) - 1
            point_clear(p)
            idx += 2
        else:
            l = int(input_data[idx+1]) - 1
            r = int(input_data[idx+2])
            out.append(str(range_sum(l, r)))
            idx += 3
            
    sys.stdout.write('\n'.join(out) + '\n')

if __name__ == '__main__':
    solve()

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

posted:
last update: