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

“学习路线”

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

前置回顾

上一篇已经得到：

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)$ 。

代码模板

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);
    }
};

完整代码

输入格式与 Luogu P3372 一致：

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

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;
}

测试用例

输入：

输出：

15
39
19

易错点

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

主练习与下一步

luogu-P3372 题解

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

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