我的算法书 - 双树状数组：区间修改与区间查询

学习路线

单点修改、区间查询 → 区间修改、单点查询 → 区间修改、区间查询（本篇）。

前置回顾

上一篇已经得到：

d_i=a_i-a_{i-1},\qquad a_x=\sum_{i=1}^{x}d_i

区间 $[l,r]$ 增加 $v$ ，等价于：

d_l\mathrel{+}=v,\qquad d_{r+1}\mathrel{-}=v

一棵树状数组足以恢复一个位置的值。本篇的新问题是：怎样恢复一整个前缀的和。

一句话算法

一棵树状数组记录差分，另一棵记录“下标乘差分”，两者组合还原原数组前缀和。

问题模型

给定长度为 $n$ 的数组 $a$ ，支持：

区间修改：给 $a_l,a_{l+1},\ldots,a_r$ 都加上 value；
区间查询：求当前 $a_l+a_{l+1}+\cdots+a_r$ 。

差分树状数组能查询单个 $a_x$ 。如果为了区间和逐个查询 $a_l,\ldots,a_r$ ，一次查询最坏仍是 $O(n\log n)$ 。

我们需要直接计算原数组的前缀和。

核心转化

如何使用差分数组 $d$ 来求原数组的区间和 $a$

设原数组前缀和为：

Pre\_sum(x)=\sum_{i=1}^{x}a_i

把 $a_i=\sum_{j=1}^{i}d_j$ 代入：

\begin{aligned} Pre\_sum(x) &=\sum_{i=1}^{x}a_i =\sum_{i=1}^{x}\sum_{j=1}^{i}d_j \end{aligned}

这一步写成展开形式更容易看清。对 $i=1$ 到 $x$ 逐行写出 $a_i$ ：

\begin{aligned} i=1&:\; a_1 = d_1 \\ i=2&:\; a_2 = d_1 + d_2 \\ i=3&:\; a_3 = d_1 + d_2 + d_3 \\ &\;\;\vdots \\ i=x&:\; a_x = d_1 + d_2 + d_3 + \cdots + d_x \end{aligned}

把上面所有行加起来，就是 $Pre\_sum(x)$ 。

每一纵列分别是 $d_1,\;d_2,\;\ldots,\;d_x$ 。按列统计：

$d_1$ 从第 $1$ 行出现到第 $x$ 行，共 $x$ 次；
$d_2$ 从第 $2$ 行出现到第 $x$ 行，共 $x-1$ 次；
一般地， $d_j$ 从第 $j$ 行出现到第 $x$ 行，共 $x - j + 1$ 次。

于是换序之后：

Pre\_sum(x) = \sum_{j=1}^{x} (x-j+1)\,d_j

再把系数拆开：

(x-j+1)\,d_j = (x+1)\,d_j - j\cdot d_j

代入并分离求和号：

\begin{aligned} Pre\_sum(x) &= \sum_{j=1}^{x}\bigl[(x+1)d_j - j\cdot d_j\bigr] \\[4pt] &= (x+1)\sum_{j=1}^{x} d_j \;-\; \sum_{j=1}^{x} j\cdot d_j \end{aligned}

这个式子只需要两种动态前缀和：

bit_diff 维护 $d_j$ ；
bit_weighted 维护 $j\cdot d_j$ 。

于是：

cpp

        1
2

        prefix_sum(x) = (x + 1) * sum(bit_diff, x)
              - sum(bit_weighted, x);

算法步骤

同步修改两棵树

对一个差分位置 $p$ 增加 $v$ 时：

cpp

        1
2

        add(bit_diff, p, v);
add(bit_weighted, p, p * v);

因此，对原数组区间 $[l,r]$ 增加 $v$ ：

cpp

        1
2
3
4
5

        add(bit_diff, l, v);
add(bit_diff, r + 1, -v);

add(bit_weighted, l, l * v);
add(bit_weighted, r + 1, -(r + 1) * v);

查询原数组前缀和

分别查询两棵树，再代入推导出的公式：

P(x)=(x+1)\operatorname{sum}_d(x)-\operatorname{sum}_{id}(x)

查询任意区间和

最后仍然使用两个前缀相减：

\operatorname{sum}(l,r)=P(r)-P(l-1)

小例子

原数组为 $[1,2,3,4,5]$ ，给 $[2,4]$ 增加 $10$ 后，差分变成：

d = [1, 11, 1, 1, -9]

查询 $P(3)$ ：

\sum_{j=1}^{3}d_j=13, \qquad \sum_{j=1}^{3}j d_j=1+22+3=26

P(3)=4\times13-26=26

这与修改后的前三项 $1+12+13=26$ 一致。

算法证明

核心不变量：bit_diff 维护当前差分 $d_i$ ，bit_weighted 维护当前加权差分 $i\cdot d_i$ 。

一次区间加只改变 $d_l$ 和 $d_{r+1}$ ；代码在两棵树中同步记录这两个变化，所以不变量保持成立。
两棵树分别正确返回
$\sum_{i=1}^{x}d_i \quad\text{和}\quad \sum_{i=1}^{x}i\cdot d_i$
将这两个结果代入恒等式
$P(x)=(x+1)\sum_{i=1}^{x}d_i-\sum_{i=1}^{x}i\cdot d_i$
得到正确的原数组前缀和。
两个正确前缀相减，得到正确的区间和。

复杂度分析

初始化：逐点加入时为 $O(n\log n)$ ；
区间修改：常数次树状数组修改， $O(\log n)$ ；
区间查询：常数次树状数组查询， $O(\log n)$ ；
空间复杂度：两棵长度为 $n$ 的树状数组，仍为 $O(n)$ 。

代码模板

C++

cpp

        1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

        #include <bits/stdc++.h>
using namespace std;

// 双树状数组：区间加、区间和。
// bit_diff 维护 b[i]，bit_weighted 维护 i * b[i]。
template <typename T>
struct RangeFenwick {
    int n = 0;
    vector<T> bit_diff, bit_weighted;

    RangeFenwick(int n = 0) { init(n); }

    void init(int size) {
        n = size;
        bit_diff.assign(n + 1, 0);
        bit_weighted.assign(n + 1, 0);
    }

    static int lowbit(int x) { return x & -x; }

    void add(vector<T> &bit, int pos, T value) {
        for (int i = pos; i <= n; i += lowbit(i)) {
            bit[i] += value;
        }
    }

    T sum(const vector<T> &bit, int pos) const {
        T answer = 0;
        for (int i = pos; i > 0; i -= lowbit(i)) {
            answer += bit[i];
        }
        return answer;
    }

    void range_add(int left, int right, T value) {
        add(bit_diff, left, value);
        add(bit_diff, right + 1, -value);
        add(bit_weighted, left, value * static_cast<T>(left));
        add(bit_weighted, right + 1, -value * static_cast<T>(right + 1));
    }

    T prefix_sum(int pos) const {
        return static_cast<T>(pos + 1) * sum(bit_diff, pos)
             - sum(bit_weighted, pos);
    }

    T range_sum(int left, int right) const {
        return prefix_sum(right) - prefix_sum(left - 1);
    }
};

Python

python

        1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

        # 双树状数组：区间加、区间和。
# bit_diff 维护差分数组 b[i]，bit_weighted 维护 i * b[i]，两者配合把区间和
# 拆成 (pos + 1) * 前缀b - 前缀(i*b)。
# 下标必须从 1 开始：lowbit(0) = 0，会让修改循环原地打转。
# Python 的 int 是任意精度，不必像 C++ 那样担心 long long 溢出。
# 调用示例：fa = RangeFenwick(5); fa.range_add(1, 3, 2); fa.range_sum(2, 5)


class RangeFenwick:
    """区间加区间和的双树状数组。

    推导：a[pos] = sum(b[1..pos])，于是
    sum(a[1..pos]) = (pos + 1) * sum(b[1..pos]) - sum(i * b[i], i = 1..pos)。
    """

    n: int
    bit_diff: list[int]
    bit_weighted: list[int]

    def __init__(self, size: int = 0) -> None:
        self.init(size)

    def init(self, size: int) -> None:
        """重置为长度 size 的空数组，方便同一个对象换一组数据复用。"""
        self.n = size
        self.bit_diff = [0] * (size + 1)
        self.bit_weighted = [0] * (size + 1)

    @staticmethod
    def lowbit(x: int) -> int:
        """块长：x 二进制最低位那个 1 代表的值，也就是 x & -x。"""
        return x & -x

    def add(self, bit: list[int], pos: int, value: int) -> None:
        """在指定的一棵树状数组 bit 上做单点加。"""
        i = pos
        while i <= self.n:
            bit[i] += value
            i += self.lowbit(i)

    def sum(self, bit: list[int], pos: int) -> int:
        """在指定的一棵树状数组 bit 上求前缀和。"""
        answer = 0
        i = pos
        while i > 0:
            answer += bit[i]
            i -= self.lowbit(i)
        return answer

    def range_add(self, left: int, right: int, value: int) -> None:
        """区间加：a[left..right] 都加上 value。

        注意 right + 1 位置的两处修改都不能漏：差分减 value，
        加权数组减 value * (right + 1)，C++ 里同样写的是右端点 + 1。
        """
        self.add(self.bit_diff, left, value)
        self.add(self.bit_diff, right + 1, -value)
        self.add(self.bit_weighted, left, value * left)
        self.add(self.bit_weighted, right + 1, -(value * (right + 1)))

    def prefix_sum(self, pos: int) -> int:
        """前缀和：a[1] + ... + a[pos]。"""
        return (pos + 1) * self.sum(self.bit_diff, pos) - self.sum(
            self.bit_weighted, pos
        )

    def range_sum(self, left: int, right: int) -> int:
        """区间和：a[left] + ... + a[right]，用两个前缀和相减得到。"""
        return self.prefix_sum(right) - self.prefix_sum(left - 1)

完整代码

输入格式与 Luogu P3372 一致：

1 l r k：给区间 $[l,r]$ 增加 $k$ ；
2 l r：查询区间 $[l,r]$ 的和。

C++

cpp

        1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

        #include <bits/stdc++.h>
using namespace std;

template <typename T>
struct RangeFenwick {
    int n = 0;
    vector<T> bit_diff, bit_weighted;

    RangeFenwick(int n = 0) { init(n); }

    void init(int size) {
        n = size;
        bit_diff.assign(n + 1, 0);
        bit_weighted.assign(n + 1, 0);
    }

    static int lowbit(int x) { return x & -x; }

    void add(vector<T> &bit, int pos, T value) {
        for (int i = pos; i <= n; i += lowbit(i)) {
            bit[i] += value;
        }
    }

    T sum(const vector<T> &bit, int pos) const {
        T answer = 0;
        for (int i = pos; i > 0; i -= lowbit(i)) {
            answer += bit[i];
        }
        return answer;
    }

    void range_add(int left, int right, T value) {
        add(bit_diff, left, value);
        add(bit_diff, right + 1, -value);
        add(bit_weighted, left, value * static_cast<T>(left));
        add(bit_weighted, right + 1, -value * static_cast<T>(right + 1));
    }

    T prefix_sum(int pos) const {
        return static_cast<T>(pos + 1) * sum(bit_diff, pos)
             - sum(bit_weighted, pos);
    }

    T range_sum(int left, int right) const {
        return prefix_sum(right) - prefix_sum(left - 1);
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m;
    cin >> n >> m;

    RangeFenwick<long long> bit(n);
    for (int i = 1; i <= n; ++i) {
        long long value;
        cin >> value;
        bit.range_add(i, i, value);
    }

    while (m--) {
        int operation, left, right;
        cin >> operation >> left >> right;
        if (operation == 1) {
            long long value;
            cin >> value;
            bit.range_add(left, right, value);
        } else {
            cout << bit.range_sum(left, right) << '\n';
        }
    }

    return 0;
}

Python

python

        1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

        # 洛谷 P3372 线段树 1（树状数组写法）：区间加 + 区间和。
# 两棵树同步维护：bit_diff 存差分 d，bit_weighted 存 i * d_i。
#   prefix_sum(x) = (x + 1) * sum(d[1..x]) - sum(i * d_i, 1..x)
# 下标从 1 开始，lowbit(0) = 0 会让循环原地打转。
# C++ 里 index * value 可能超过 int，用 long long；Python int 任意精度，不必特判。

import sys


class RangeFenwick:
    """tree[i] 维护对应数组在区间 (i - lowbit(i), i] 上的和。"""

    def __init__(self, size: int = 0) -> None:
        self.n = size
        self.bit_diff = [0] * (size + 1)
        self.bit_weighted = [0] * (size + 1)

    @staticmethod
    def lowbit(x: int) -> int:
        # 最低位那个 1 代表的块长，即 x & -x。
        return x & -x

    def add(self, bit: list[int], pos: int, value: int) -> None:
        i = pos
        while i <= self.n:
            bit[i] += value
            i += self.lowbit(i)

    def sum(self, bit: list[int], pos: int) -> int:
        answer = 0
        i = pos
        while i > 0:
            answer += bit[i]
            i -= self.lowbit(i)
        return answer

    def range_add(self, left: int, right: int, value: int) -> None:
        # 差分：左端点 +value，右端点后一位 -value。
        self.add(self.bit_diff, left, value)
        self.add(self.bit_diff, right + 1, -value)
        # 加权树同步：位置 i 的增量是 i * d_i，右端点后一位用 r + 1。
        self.add(self.bit_weighted, left, value * left)
        self.add(self.bit_weighted, right + 1, -value * (right + 1))

    def prefix_sum(self, pos: int) -> int:
        return (pos + 1) * self.sum(self.bit_diff, pos) - self.sum(self.bit_weighted, pos)

    def range_sum(self, left: int, right: int) -> int:
        return self.prefix_sum(right) - self.prefix_sum(left - 1)


def main() -> None:
    data = sys.stdin.buffer.read().split()
    if not data:
        return
    it = iter(map(int, data))

    n = next(it)
    m = next(it)

    bit = RangeFenwick(n)
    for i in range(1, n + 1):
        # 用区间加 (i, i) 建树，与 C++ 逐点写入的做法一致。
        bit.range_add(i, i, next(it))

    out: list[str] = []
    for _ in range(m):
        operation = next(it)
        left = next(it)
        right = next(it)
        if operation == 1:
            value = next(it)
            bit.range_add(left, right, value)
        else:
            out.append(str(bit.range_sum(left, right)))

    if out:
        sys.stdout.write("\n".join(out) + "\n")


if __name__ == "__main__":
    main()

测试用例

输入：

输出：

15
39
19

易错点

必须明确第二棵树维护的是 $i\cdot d_i$ ，因此公式对应 (x + 1)；不要与维护 $(i-1)d_i$ 的另一种写法混用。
修改右边界时，下标和权值都要使用 r + 1。
两棵树必须同步修改；漏改任何一棵都会破坏前缀和公式。
乘法 index * value 容易超过 int，应使用 long long。
这个方法依赖“区间加”和“区间和”的线性结构，不能直接处理区间赋值、区间乘法或一般区间最值。

主练习与下一步

luogu-P3372 题解

P3372 的操作模型正是区间加与区间和。题目名称虽然写着“线段树”，但在只有加法和求和时，双树状数组同样适用；题解中给出了两种做法。

求和主线到这里结束。选学分支：树状数组维护前缀最值。