Official

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

GPT 5.2 High

概要

区間加算(操作1)と区間和(操作2)を \(Q\) 回処理する問題です。各操作を高速に処理するために、Fenwick Tree(BIT)を工夫して使います。

考察

素朴に操作1のたびに \(l..r\) を全て更新すると、1回あたり最悪 \(O(N)\) かかり、全体で \(O(NQ)\) となって \(N,Q \le 10^5\) では間に合いません。

必要なのは次の2つをどちらも高速に行うことです。

  • 区間 \([l,r]\) への加算(range add)
  • 区間 \([l,r]\) の総和(range sum)

BITは通常「点加算・前計算(prefix sum)」が得意ですが、差分配列や累積和の関係を利用すると「区間加算・区間和」も \(O(\log N)\) で処理できます。

ポイントは、 - 「区間加算」は差分(端点)にだけ更新を入れる - 「区間和」は“各要素の値”をさらに累積して求める必要があるため、BITを2本使って式を整える
という発想です。

アルゴリズム

1) 区間加算を差分で表す

配列 \(A\) に対し、差分配列 \(D\) を - \(D[1] = A[1]\) - \(D[i] = A[i] - A[i-1] \ (i\ge2)\) とします。

区間 \([l,r]\)\(v\) を足す操作は、差分では - \(D[l] += v\) - \(D[r+1] -= v\)(ただし \(r+1 \le N\) のとき) の2点更新で済みます。

すると各要素は - \(A[x] = \sum_{i=1}^{x} D[i]\) で復元できます。

BIT(bit1)でこの差分 \(D\) の「点加算・prefix sum」を管理すると、任意の \(A[x]\) は - \(A[x] = \text{bit1.sum}(x)\) で求まります。

2) 区間和は「prefix sum のさらに累積」

区間和を取るには prefix sum \(S(x)=\sum_{i=1}^{x} A[i]\) が欲しいです。

差分を使うと \( A[i]=\sum_{j=1}^{i}D[j] \) なので \( S(x)=\sum_{i=1}^{x}A[i] =\sum_{i=1}^{x}\sum_{j=1}^{i}D[j] =\sum_{j=1}^{x}D[j]\cdot(x-j+1) \) これを整理すると \( S(x)= (x+1)\sum_{j=1}^{x}D[j] - \sum_{j=1}^{x}D[j]\cdot j \)

ここで - \(\sum_{j=1}^{x} D[j]\) は bit1.sum(x) - \(\sum_{j=1}^{x} D[j]\cdot j\) を管理するために、もう1本BIT(bit2)を用意して \(D[j]\cdot j\) のprefix sumを取れるようにする

実装上は、よく使われる等価な形として \( S(x)= \text{bit1.sum}(x)\cdot x - \text{bit2.sum}(x) \) となるように更新値を設計します。

3) 更新(range add)の具体的なBIT2更新

区間 \([l,r]\)\(v\) を足すと差分は - \(D[l] += v\) - \(D[r+1] -= v\)

このとき、bit2 には「式を成立させるための補正値」を入れます(コードの通り):

  • bit1.add(l, v), bit1.add(r+1, -v)
  • bit2.add(l, v*(l-1)), bit2.add(r+1, -v*r)

すると - prefix_sum(x) = bit1.sum(x)*x - bit2.sum(x) が「\(1..x\) の総和」になります。

最後に区間和は \( \sum_{i=l}^{r}A[i] = S(r)-S(l-1) \) で求まります。

4) 初期配列 \(C_i\) の反映

初期値も「区間 \([i,i]\)\(C_i\) を加算」とみなして同じ range_add を呼べば統一的に処理できます(コードではこれを採用)。

計算量

  • 時間計算量: \(O((N+Q)\log N)\)
    (初期化 \(N\) 回 + 各クエリ \(O(\log N)\)
  • 空間計算量: \(O(N)\)
    (BITを2本持つ)

実装のポイント

  • 添字は1-indexedで実装する(BITと相性が良い)。入力も \(1..N\) なのでそのまま扱える。

  • 更新で \(r+1 \le N\) のときだけ \(r+1\) へ反映する(範囲外アクセス防止)。

  • 出力が多いので、sys.stdin.buffer.read() と配列化・最後にまとめて出力することで高速化している。

  • 値の合計は最大で大きくなるため(\(10^5 \times 10^4 \times 10^5\) 規模もあり得る)、Pythonの int(任意精度)で安全に扱える。

    ソースコード

import sys

class BIT:
    __slots__ = ("n", "bit")
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)

    def add(self, i, x):
        n = self.n
        while i <= n:
            self.bit[i] += x
            i += i & -i

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.bit[i]
            i -= i & -i
        return s

def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    it = iter(data)
    N = next(it)
    Q = next(it)

    bit1 = BIT(N)
    bit2 = BIT(N)

    def range_add(l, r, v):
        bit1.add(l, v)
        if r + 1 <= N:
            bit1.add(r + 1, -v)
        bit2.add(l, v * (l - 1))
        if r + 1 <= N:
            bit2.add(r + 1, -v * r)

    def prefix_sum(x):
        return bit1.sum(x) * x - bit2.sum(x)

    def range_sum(l, r):
        return prefix_sum(r) - prefix_sum(l - 1)

    # initialize with C_i
    for i in range(1, N + 1):
        c = next(it)
        if c:
            range_add(i, i, c)

    out = []
    for _ in range(Q):
        t = next(it)
        if t == 1:
            l = next(it); r = next(it); v = next(it)
            range_add(l, r, v)
        else:
            l = next(it); r = next(it)
            out.append(str(range_sum(l, r)))

    sys.stdout.write("\n".join(out))

if __name__ == "__main__":
    main()

この解説は gpt-5.2-high によって生成されました。

posted:
last update: