我的算法书 - 树状数组与差分：区间修改与单点查询

学习路线

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

前置回顾

本篇只组合两个已经学过的结论：

基础树状数组支持差分数组的单点修改和前缀查询；
差分满足
$d_i=a_i-a_{i-1},\qquad a_i=\sum_{j=1}^{i}d_j$

不再重复推导 lowbit 和普通差分，只关注两者怎样组合成在线算法。

一句话算法

把区间加法记在差分数组的两个边界上，再用树状数组求差分前缀和恢复单点值。

问题模型

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

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

如果每次区间修改都逐项更新，最坏需要 $O(n)$ 。普通差分虽然能 $O(1)$ 记录修改，但通常要等所有修改结束后再统一还原，不能高效处理修改与查询交替出现的情况。

核心转化

令：

d_i=a_i-a_{i-1},\qquad a_0=0

给区间 $[l,r]$ 中每个数增加 $v$ ，只会改变差分数组的两个边界：

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

任意位置 $x$ 的值等于差分前缀和：

a_x=d_1+d_2+\cdots+d_x

因此，问题变成基础树状数组最擅长的模型：

原问题：区间修改 + 单点查询
             |
             v
差分数组：两个单点修改 + 一次前缀查询

算法步骤

初始化

读入原数组时计算相邻差：

cpp

        1
2
3
4
5
6

        previous = 0;
for (int i = 1; i <= n; ++i) {
    cin >> value;
    bit.add(i, value - previous);
    previous = value;
}

树状数组实际维护的是 $d$ ，不是原数组 $a$ 。

区间修改

对 $[l,r]$ 增加 $v$ ：

cpp

        1
2

        add(l, v);
if (r < n) add(r + 1, -v);

d[l] += v 表示影响从 $l$ 开始，d[r+1] -= v 表示影响在 $r$ 之后停止。

单点查询

查询 $a_x$ ：

cpp

        1

        point_query(x) = prefix_sum(x);

例如原数组为 $[1,2,3,4,5]$ ，给 $[2,4]$ 加 $10$ ：

原差分：  1  1  1  1  1
边界修改：   +10      -10
新差分：  1 11  1  1 -9
前缀还原：1 12 13 14  5

查询位置 $3$ 时，差分前缀和为 $1+11+1=13$ 。

算法证明

核心不变量：树状数组始终维护当前数组的差分 $d$ ，且

a_x=\sum_{i=1}^{x}d_i

区间 $[l,r]$ 增加 $v$ 后：

$d_l$ 增加 $v$ ，所以从 $l$ 开始的前缀和都多出 $v$ ；
$d_{r+1}$ 减少 $v$ ，所以从 $r+1$ 开始抵消这次增加；
因而只有 $a_l,\ldots,a_r$ 增加 $v$ 。

树状数组正确完成两个差分点修改，并正确返回差分前缀和，所以每次单点查询得到当前真实的 $a_x$ 。

复杂度分析

初始化： $O(n\log n)$ ；
区间修改：两次单点修改， $O(\log n)$ ；
单点查询：一次前缀查询， $O(\log n)$ ；
空间复杂度： $O(n)$ 。

代码模板

模板中的 add 操作修改一个差分位置，日常使用的接口是 range_add 和 point_query：

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

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

// 维护差分数组：区间加、单点查询。
// 下标必须从 1 开始。
template <typename T>
struct RangeAddPointQueryFenwick {
    int n = 0;
    vector<T> tree;

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

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

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

    // 给差分数组的一个位置增加 value。
    void add(int pos, T value) {
        for (int i = pos; i <= n; i += lowbit(i)) {
            tree[i] += value;
        }
    }

    T prefix_sum(int pos) const {
        T answer = 0;
        for (int i = pos; i > 0; i -= lowbit(i)) {
            answer += tree[i];
        }
        return answer;
    }

    void range_add(int left, int right, T value) {
        add(left, value);
        if (right + 1 <= n) add(right + 1, -value);
    }

    T point_query(int pos) const {
        return prefix_sum(pos);
    }
};

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

        # 区间加（维护差分数组）、单点查询。
# 下标必须从 1 开始：lowbit(0) = 0，会让修改循环原地打转。
# Python 的 int 是任意精度，不必像 C++ 那样担心 long long 溢出。
# 调用示例：fa = RangeAddPointQueryFenwick(5); fa.range_add(1, 3, 2); fa.point_query(3)


class RangeAddPointQueryFenwick:
    """树状数组上维护差分：tree[i] 维护 diff[i - lowbit(i) + 1 .. i] 的元素和。

    区间加等价于差分数组上两次单点加，单点查询等价于差分前缀和。
    """

    n: int
    tree: list[int]

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

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

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

    def add(self, pos: int, value: int) -> None:
        """给差分数组的一个位置增加 value：diff[pos] += value。"""
        i = pos
        while i <= self.n:
            self.tree[i] += value
            i += self.lowbit(i)

    def prefix_sum(self, pos: int) -> int:
        """差分前缀和：diff[1] + ... + diff[pos]，恰好等于 a[pos] 的当前值。"""
        answer = 0
        i = pos
        while i > 0:
            answer += self.tree[i]
            i -= self.lowbit(i)
        return answer

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

        等价于 diff[left] += value 且 diff[right + 1] -= value；
        right + 1 越过末尾时无需修改（C++ 同样有这个 if 保护）。
        """
        self.add(left, value)
        if right + 1 <= self.n:
            self.add(right + 1, -value)

    def point_query(self, pos: int) -> int:
        """单点查询：返回 a[pos] 的当前值。"""
        return self.prefix_sum(pos)

完整代码

输入格式与 Luogu P3368 一致：

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

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

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

template <typename T>
struct RangeAddPointQueryFenwick {
    int n = 0;
    vector<T> tree;

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

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

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

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

    T prefix_sum(int pos) const {
        T answer = 0;
        for (int i = pos; i > 0; i -= lowbit(i)) {
            answer += tree[i];
        }
        return answer;
    }

    void range_add(int left, int right, T value) {
        add(left, value);
        if (right < n) add(right + 1, -value);
    }

    T point_query(int pos) const {
        return prefix_sum(pos);
    }
};

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

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

    RangeAddPointQueryFenwick<long long> bit(n);
    long long previous = 0;
    for (int i = 1; i <= n; ++i) {
        long long value;
        cin >> value;
        bit.add(i, value - previous);
        previous = value;
    }

    while (m--) {
        int operation;
        cin >> operation;
        if (operation == 1) {
            int left, right;
            long long value;
            cin >> left >> right >> value;
            bit.range_add(left, right, value);
        } else {
            int pos;
            cin >> pos;
            cout << bit.point_query(pos) << '\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

        # 洛谷 P3368 树状数组 2：区间加 + 单点查。
# 树状数组维护差分数组 d，point_query(x) = sum(d[1..x]) 才是原数组 a[x]。
# 下标从 1 开始，lowbit(0) = 0 会让循环原地打转。
# C++ 用 long long 防溢出；Python int 是任意精度，不必特判。

import sys


class RangeAddPointQueryFenwick:
    """tree[i] 维护差分数组 d[i - lowbit(i) + 1 .. i] 的元素和。"""

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

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

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

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

    def range_add(self, left: int, right: int, value: int) -> None:
        self.add(left, value)
        # 只有 right < n 时 right + 1 才是合法差分下标；等于 n + 1 时 add 循环不进入，可省略。
        if right < self.n:
            self.add(right + 1, -value)

    def point_query(self, pos: int) -> int:
        return self.prefix_sum(pos)


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 = RangeAddPointQueryFenwick(n)
    previous = 0
    for i in range(1, n + 1):
        value = next(it)
        # 初始化喂的是差分 a[i] - a[i-1]，不是 a[i] 本身。
        bit.add(i, value - previous)
        previous = value

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

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


if __name__ == "__main__":
    main()

测试用例

输入：

输出：

3
13
5

位置 $5$ 不在修改区间内，所以最后一次查询仍为 $5$ 。

易错点

树状数组维护的是差分数组，point_query(x) 才是原数组的 $a_x$ 。
只有当 r < n 时，r+1 才是有效差分下标。
初始化时要加入 a[i] - a[i-1]，不能直接加入 a[i]。
区间修改的右边界必须写成 -value，用于停止影响。

主练习与下一步

luogu-P3368 题解

这道题直接检验“区间加转成两个差分点修改，单点值转成差分前缀和”。

下一篇：双树状数组：区间修改与区间查询。