Official

E - 花壇の水やり管理 / Flowerbed Watering Management Editorial by admin Editorial for printing

Claude 4.6 Opus (Thinking)

概要

横一列に並んだ \(N\) 個の花壇に対して、「区間への一様加算」と「区間の合計値の取得」を効率的に処理する問題です。Binary Indexed Tree (BIT) を2本使うことで、両方の操作を \(O(\log N)\) で実現します。

考察

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

最も単純な方法は、配列をそのまま管理することです。

  • 操作1(区間加算): \(l\) から \(r\) までの各要素に \(v\) を加算 → \(O(N)\)
  • 操作2(区間合計): \(l\) から \(r\) までの各要素を足し合わせる → \(O(N)\)

\(N\) と \(Q\) がともに最大 \(10^5\) なので、最悪 \(O(NQ) = O(10^{10})\) となり、制限時間に間に合いません。

重要な気づき

この問題は「区間加算」と「区間和取得」の2つの操作を高速に行う必要があります。これは典型的なデータ構造の問題で、以下のような手法で解決できます:

  • セグメント木(遅延伝播付き)
  • BIT(Binary Indexed Tree)を2本使う方法

ここでは、実装が比較的シンプルな BIT を2本使う方法 を採用します。

アルゴリズム

BIT による区間加算・区間和取得の原理

区間 \([l, r]\) に値 \(x\) を加算する操作を、累積和の観点で考えます。

差分配列 \(d\) を導入し、\(d[i]\) は「位置 \(i\) に加算された増分」とすると、位置 \(1\) から \(i\) までの合計(prefix sum)は:

\[\text{prefix\_sum}(i) = \sum_{j=1}^{i} \sum_{k=1}^{j} d[k]\]

これを整理すると:

\[\text{prefix\_sum}(i) = \sum_{k=1}^{i} d[k] \cdot (i - k + 1) = (i+1)\sum_{k=1}^{i} d[k] - \sum_{k=1}^{i} d[k] \cdot k\]

この式には2種類の累積和が登場します:

  • \(\text{bit1}\):\(d[k]\) の累積和を管理
  • \(\text{bit2}\):\(d[k] \cdot (k-1)\) の累積和を管理

これにより:

\[\text{prefix\_sum}(i) = \text{bit1\_sum}(i) \times i - \text{bit2\_sum}(i)\]

区間 \([l, r]\) への加算は、\(\text{bit1}\) と \(\text{bit2}\) それぞれに対して2箇所ずつ更新するだけで済みます。

初期値の処理

初期値 \(C_i\) については、操作による加算分とは分離して、あらかじめ累積和配列 prefix を作っておきます。操作2の回答時に、初期値の区間和と BIT で管理している加算分の区間和を足し合わせます。

具体例

\(N=5\)、\(C = [3, 1, 4, 1, 5]\) のとき: - prefix = [0, 3, 4, 8, 9, 14] - 操作 1 2 4 2(花壇2〜4に2リットル追加)→ BIT に区間加算 - 操作 2 1 5(花壇1〜5の合計)→ \(\text{prefix}[5] - \text{prefix}[0] + \text{bit.range\_sum}(1,5) = 14 + 6 = 20\)

計算量

  • 時間計算量: \(O((N + Q) \log N)\)(初期化に \(O(N)\)、各クエリに \(O(\log N)\))
  • 空間計算量: \(O(N)\)(BIT 2本と累積和配列)

実装のポイント

  • 初期値と加算分の分離:初期値は累積和配列で前処理し、BIT は操作1による加算分のみを管理します。これにより BIT の初期化が不要になり、実装がシンプルになります。

  • 1-indexed:BIT は 1-indexed で扱うのが自然です。入力の花壇番号もそのまま使えます。

  • 出力の一括処理:sys.stdout.write でまとめて出力することで、Python での I/O のオーバーヘッドを軽減しています。

    ソースコード

import sys
input = sys.stdin.readline

class BIT:
    def __init__(self, n):
        self.n = n
        self.bit1 = [0] * (n + 1)
        self.bit2 = [0] * (n + 1)
    
    def _add(self, bit, i, x):
        while i <= self.n:
            bit[i] += x
            i += i & (-i)
    
    def _sum(self, bit, i):
        s = 0
        while i > 0:
            s += bit[i]
            i -= i & (-i)
        return s
    
    def range_add(self, l, r, x):
        self._add(self.bit1, l, x)
        self._add(self.bit1, r + 1, -x)
        self._add(self.bit2, l, x * (l - 1))
        self._add(self.bit2, r + 1, -x * r)
    
    def prefix_sum(self, i):
        return self._sum(self.bit1, i) * i - self._sum(self.bit2, i)
    
    def range_sum(self, l, r):
        return self.prefix_sum(r) - self.prefix_sum(l - 1)

def main():
    N, Q = map(int, input().split())
    C = list(map(int, input().split()))
    
    bit = BIT(N)
    
    # Initialize: add each C[i] as a point update (range of length 1)
    # Use prefix sum array for initial values
    prefix = [0] * (N + 1)
    for i in range(N):
        prefix[i + 1] = prefix[i] + C[i]
    
    out = []
    for _ in range(Q):
        query = list(map(int, input().split()))
        if query[0] == 1:
            l, r, v = query[1], query[2], query[3]
            bit.range_add(l, r, v)
        else:
            l, r = query[1], query[2]
            ans = (prefix[r] - prefix[l - 1]) + bit.range_sum(l, r)
            out.append(ans)
    
    sys.stdout.write('\n'.join(map(str, out)) + '\n')

main()

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

posted:
last update: