算法数据结构

线段树和树状数组：单点修改、区间和查询

字数统计: 3.1k字   |   阅读时长: 15分
, ,

摘要: 线段树和树状数组，单点修改与区间查询，原理与实现

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

本文通以 307. 区域和检索 - 数组可修改 作为模板题来看一下线段树和树状数组的原理与实现。

本文的树状数组和线段树是支持单点修改和区间查询这一套需求的。如果需要区间修改，参考文章 线段树：区间修改、区间查询，区间最值问题。

模板题：307. 区域和检索 - 数组可修改

给你一个数组 nums ，请你完成两类查询。

  1. 其中一类查询要求 更新 数组 nums 下标对应的值
  2. 另一类查询要求返回数组 nums 中索引 left 和索引 right 之间（ 包含 ）的nums元素的 和 ，其中 left <= right

实现 NumArray 类：

  • NumArray(int[] nums) 用整数数组 nums 初始化对象
  • void update(int index, int val) 将 nums[index] 的值 更新 为 val

提示：

1
2
3
4
5
6
1 <= nums.length <= 3e4
-100 <= nums[i] <= 100
0 <= index < nums.length
-100 <= val <= 100
0 <= left <= right < nums.length
调用 update 和 sumRange 方法次数不大于 3e4

  • int sumRange(int left, int right) 返回数组 nums 中索引 left 和索引 right 之间（ 包含 ）的nums元素的 和 （即，nums[left] + nums[left + 1], …, nums[right]）

示例 1：
输入：
[“NumArray”, “sumRange”, “update”, “sumRange”]
[[[1, 3, 5]], [0, 2], [1, 2], [0, 2]]
输出：
[null, 9, null, 8]
解释：
NumArray numArray = new NumArray([1, 3, 5]);
numArray.sumRange(0, 2); // 返回 1 + 3 + 5 = 9
numArray.update(1, 2); // nums = [1,2,5]
numArray.sumRange(0, 2); // 返回 1 + 2 + 5 = 8

算法1 : 线段树

基于链表

线段树是一个查询和修改复杂度都为$\log n$的数据结构。主要用于数组的单点修改、单点查询、区间修改，区间查询。

本题需要支持的操作是单点修改和区间查询。

  • 建树 build(int start, int end, vector<int> vals);
  • 单点更新 update(int index, int val);
  • 区间查询 query(int start, int end);

线段树是一种平衡二叉树(树高为$\log N$)，的每个节点代表一个区间，保存区间的两个端点和区间的和这三个值。

线段树节点的定义：

1
2
3
4
5
6
7
8
9
10
11
struct STNode {
    int start;
    int end;
    int sum;
    STNode *left;
    STNode *right;
    STNode(){}
    STNode(int start, int end, int sum, STNode *left=nullptr, STNode *right=nullptr)
        :start(start),end(end),sum(sum),left(left),right(right){}
    ~STNode(){}
};

start 和 end 是区间的端点，sum 是区间的和(也可以是最大，最小，异或，模质数下乘法)。

叶子节点表示的就是数组的元素本身，sum 是数组元素值，start = end 为元素在数组中的下标。

一个非叶子节点表示区间 [left, right]，它的两个子节点表示区间 [left, mid], [mid+1, right], mid = (left + right) / 2，且节点的 sum 值为左右两个子节点的 sum 值的和。

根节点表示整个数组的范围。

例如数组 nums = [2, 1, 5, 3, 4]，建的线段树如下图所示：根节点的区间就是整个区间 [0-4], 值是区间所有元素的和 15，各个子节点表示各自的子区间以及子区间上的和，直到叶子节点表示数组元素本身。

<1> 由给定数组建树

与一般的平衡树差不多，递归建树，不同的是叶子节点的判断是根据 start == end。两个子节点建好后，再建当前节点，当前节点的 sum 值需要用到子节点的 sum 值。

1
2
3
4
5
6
7
8
9
STNode* _build(int start, int end, const vector<int>& vals)
{
    if(start == end)
        return new STNode(start, end, vals[start]);
    int mid = start + (end - start) / 2;
    STNode *left = _build(start, mid, vals);
    STNode *right = _build(mid + 1, end, vals);
    return new STNode(start, end, left -> sum + right -> sum, left, right);
}

<2> 单点更新

通过 index 与 mid = (left + right) / 2 的关系找到 index 在哪个子树：然后沿着子树链条往下走直到叶子，更新叶子的 sum 值。回溯之后更新当前节点的 sum 值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
void _point_update(STNode *root, int index, int val)
{
    if(root -> start == root -> end && root -> end == index)
    {
        root -> sum = val;
        return;
    }
    int mid = root -> start + (root -> end - root -> start) / 2;
    if(index <= mid) // 更新位置在左子树
        _point_update(root -> left, index, val);
    else
        _point_update(root -> right, index, val);
    root -> sum = root -> left -> sum + root -> right -> sum;
}

<3> 区间查询

1
2
3
4
5
6
7
8
9
10
11
12
int _range_query(STNode *root, int i, int j)
{
    if(root -> start == i && root -> end == j)
        return root -> sum;
    int mid = root -> start + (root -> end - root -> start) / 2;
    if(j <= mid) // 查询的区间在左子树
        return _range_query(root -> left, i, j);
    else if(i > mid) // 查询的区间在右子树
        return _range_query(root -> right, i, j);
    else
        return _range_query(root -> left, i, mid) + _range_query(root -> right, mid + 1, j);
}

代码(c++，模板)

单点更新，区间查询的模板写法，没有增删。

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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
struct STNode {
    int start;
    int end;
    int sum;
    STNode *left;
    STNode *right;
    STNode(){}
    STNode(int start, int end, int sum)
        :start(start),end(end),sum(sum),left(nullptr),right(nullptr){}
    STNode(int start, int end, int sum, STNode *left, STNode *right)
        :start(start),end(end),sum(sum),left(left),right(right){}
};

class SegmentTree {
public:
    SegmentTree(){
        root = nullptr;
    }

    ~SegmentTree()
    {
        if(root)
        {
            delete_sub_tree(root);
        }
    }

    void build(int start, int end, const vector<int>& vals)
    {
        if(end >= start)
            root = _build(start, end, vals);
    }

    void point_update(int index, int val)
    {
        _point_update(root, index, val);
    }

    int range_query(int i, int j)
    {
        return _range_query(root, i, j);
    }

private:
    STNode *root;

    void delete_sub_tree(const STNode* node) 
    {
        if(node)
        {
            if(node -> left)
                delete_sub_tree(node -> left);
            if(node -> right)
                delete_sub_tree(node -> right);
        }
        delete node;
        node = nullptr;
    }

    STNode* _build(int start, int end, const vector<int>& vals)
    {
        if(start == end)
            return new STNode(start, end, vals[start]);
        int mid = start + (end - start) / 2;
        STNode *left = _build(start, mid, vals);
        STNode *right = _build(mid + 1, end, vals);
        return new STNode(start, end, left -> sum + right -> sum, left, right);
    }

    void _point_update(STNode *root, int index, int val)
    {
        if(root -> start == root -> end && root -> end == index)
        {
            root -> sum = val;
            return;
        }
        int mid = root -> start + (root -> end - root -> start) / 2;
        if(index <= mid) // 更新位置在左子树
            _point_update(root -> left, index, val);
        else
            _point_update(root -> right, index, val);
        root -> sum = root -> left -> sum + root -> right -> sum;
    }

    int _range_query(STNode *root, int i, int j)
    {
        if(root -> start == i && root -> end == j)
            return root -> sum;
        int mid = root -> start + (root -> end - root -> start) / 2;
        if(j <= mid) // 查询的区间在左子树
            return _range_query(root -> left, i, j);
        else if(i > mid) // 查询的区间在右子树
            return _range_query(root -> right, i, j);
        else
            return _range_query(root -> left, i, mid) + _range_query(root -> right, mid + 1, j);
    }
};

class NumArray {
public:
    NumArray(vector<int>& nums) {
        int n = nums.size();
        int start = 0, end = n - 1;
        sttree = SegmentTree();
        sttree.build(start, end, nums);
    }

    void update(int i, int val) {
        sttree.point_update(i, val);
    }

    int sumRange(int i, int j) {
        return sttree.range_query(i, j);
    }

private:
    SegmentTree sttree;
};

基于数组

数据数组 nums 的范围是 [0..n]，则线段树的根节点表示的区间总是 [0, n - 1]。

因为线段树是几乎被填满的二叉树，因此数组的实现更省空间。线段树的数据存在数组 st_vec 中。st_vec[i] 表示节点 i 的值，这里是和，也可以是最值等。

节点 i 的两个子节点分别为左子节点 i*2 右子节点 i*2+1。线段树数据从下标 1 开始，1 为根节点。

以 nums 建树时，st_vec 开 4*n 的空间是够的。

建树，查询，更新的接口如下：

1
2
3
void _build(int node, int nodeLeft, int nodeRight, const vector<int>& arr)
int _range_query(int node, int nodeLeft, int nodeRight, int start, int end)
void _point_update(int node, int nodeLeft, int nodeRight, int idx, int val)

数组形式的线段树，节点 node 表示的区间 [nodeLeft, nodeRight] 是在递归中作为参数带着的。而链式写法中，这个信息是保存在节点上的，递归中不带该参数。

代码(c++, 模板)

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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
class SeqSegmentTree
{
public:
    SeqSegmentTree()
    {
        st_vec = vector<int>();
        n = -1;
    }

    void build(const vector<int>& arr)
    {
        n = arr.size();
        int start = 0, end = n - 1;
        if(start <= end)
        {
            st_vec.resize(n * 4);
            _build(1, 0, n - 1, arr);
        }
    }

    int range_query(int start, int end)
    {
        // 这里 [0, n - 1] 跟着参数走
        // 链式写法中保存在节点
        return _range_query(1, 0, n - 1, start, end);
    }

    void point_update(int idx, int val)
    {
        _point_update(1, 0, n - 1, idx, val);
    }

private:
    vector<int> st_vec;
    int n; // 保存根节点的范围 [0, n - 1]

    void _point_update(int node, int nodeLeft, int nodeRight, int idx, int val)
    {
        // [nodeLeft, nodeRight] 与 idx 的关系
        if(nodeLeft == nodeRight && nodeLeft == idx)
        {
            st_vec[node] = val;
            return;
        }

        int nodeMid = (nodeLeft + nodeRight) / 2;
        int left_son = node * 2, right_son = node * 2 + 1;
        if(idx <= nodeMid)
            _point_update(left_son, nodeLeft, nodeMid, idx, val);
        else
            _point_update(right_son, nodeMid + 1, nodeRight, idx, val);
        st_vec[node] = st_vec[left_son] + st_vec[right_son];
    }

    void _build(int node, int nodeLeft, int nodeRight, const vector<int>& arr)
    {
        // node 表示 [nodeLeft, nodeRight]
        if(nodeLeft == nodeRight)
        {
            st_vec[node] = arr[nodeLeft];
            return;
        }

        int nodeMid = (nodeLeft + nodeRight) / 2;
        int left_son = node * 2, right_son = node * 2 + 1;
        _build(left_son, nodeLeft, nodeMid, arr);
        _build(right_son, nodeMid + 1, nodeRight, arr);
        st_vec[node] = st_vec[left_son] + st_vec[right_son];
    }

    int _range_query(int node, int nodeLeft, int nodeRight, int start, int end)
    {
        // [nodeLeft, nodeRight] 与 [start, end] 的关系
        // 得出交集，返回和
        if(nodeLeft == start && nodeRight == end)
            return st_vec[node];

        int nodeMid = (nodeLeft + nodeRight) / 2;
        int left_son = node * 2, right_son = node * 2 + 1;

        if(end <= nodeMid)
            return _range_query(left_son, nodeLeft, nodeMid, start, end);
        else if(start > nodeMid)
            return _range_query(right_son, nodeMid + 1, nodeRight, start, end);
        else // [nodeLeft, start, nodeMid, end, nodeRight]
            return _range_query(left_son, nodeLeft, nodeMid, start, nodeMid)
                + _range_query(right_son, nodeMid + 1, nodeRight, nodeMid + 1, end);
    }
};

class NumArray {
public:
    NumArray(vector<int>& nums) {
        seqsttree = SeqSegmentTree();
        seqsttree.build(nums);
    }

    void update(int i, int val) {
        seqsttree.point_update(i, val);
    }

    int sumRange(int i, int j) {
        return seqsttree.range_query(i, j);
    }

private:
    SeqSegmentTree seqsttree;
};

算法2: 树状数组

树状数组只能维护前缀信息(本题是前缀和)，区间和可以用两个前缀和的差得到，所以区间和的问题用树状数组可做。

例如一个长度为 8 的数组 a ，用树状数组 b 表示的前缀和结构如下图：

树状数组中的值与原数组值的关系如下：

b[1] = a[1]
b[2] = a[1] + a[2]
b[3] = a[3]
b[4] = a[1] + a[2] + a[3] + a[4]
b[5] = a[5]
b[6] = a[5] + a[6]
b[7] = a[7]
b[8] = a[1] + a[2] + a[3] + a[4] + a[5] + a[6] + a[7] + a[8]

对 b[index] 更新(增加delta)后，会有比 index 大的某些位置也会受影响，例如更新 b[5] 之后，会依次影响到 b[6] 和 b[8]，参考图里画的树。下一个受影响的的位置是 index 加上它最低位的 1 表示的数。最低位的1表示的数，实现方法是 x & (-x)。

1
2
3
4
int _lowbit(int x)
{
    return x & (-x);
}

<1> 单点更新

index 上的值增加 delta 后，依次把后续受影响的位置也加上 delta。

1
2
3
4
5
6
7
8
9
void update(int index, int delta)
{
    int n = sums.size();
    while(index < n)
    {
        sums[index] += delta;
        index += _lowbit(index);
    }
}

<2> 单点查询

查询前缀和。例如：

b[4]=a[1]+a[2]+a[3]+a[4];
b[5]=a[5];
可以推出: query(5) = b[4]+b[5];
序号写为二进制: query(101) = b[(100)] + b[(101)];
第一次 101,减去最低位的 1 就是 100;

也就是单点更新的逆操作：

1
2
3
4
5
6
7
8
9
10
int query(int i)
{
    int sum = 0;
    while(i > 0)
    {
        sum += sums[i];
        i -= _lowbit(i);
    }
    return sum;
}

区间查询可以通过两次单点查询得到：

1
2
// 求 [i..j] 的区间和
query(j + 1) - query(i);

代码(c++，模板)

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
class BIT {
public:
    BIT():sums(1, 0){}
    BIT(int n):sums(n + 1, 0){}

    void update(int index, int delta)
    {
        int n = sums.size();
        while(index < n)
        {
            sums[index] += delta;
            index += _lowbit(index);
        }
    }

    int query(int i)
    {
        int sum = 0;
        while(i > 0)
        {
            sum += sums[i];
            i -= _lowbit(i);
        }
        return sum;
    }

private:
    vector<int> sums;

    int _lowbit(int x)
    {
        return x & (-x);
    }
};

class NumArray_2 {
public:
    // vec初始化完后 nums.size() 就变为0了
    NumArray_2(vector<int> nums):vec(move(nums)),bit(vec.size()) {
        int n = vec.size();
        for(int i = 0; i < n; ++i)
            bit.update(i + 1, vec[i]);
    }

    void update(int i, int val) {
        bit.update(i + 1, val - vec[i]);
        vec[i] = val;
    }

    int sumRange(int i, int j) {
        return bit.query(j + 1) - bit.query(i);
    }

private:
    vector<int> vec;
    BIT bit;
};

线段树和树状数组的区别和联系

具体区别和联系如下：

  1. 时间复杂度相同, 但是树状数组的常数优于线段树。

  2. 树状数组的作用被线段树完全涵盖, 凡是可以使用树状数组解决的问题, 使用线段树一定可以解决, 但是线段树能够解决的问题树状数组未必能够解决。

  3. 树状数组的代码量比线段树小很多。

  4. 线段树和树状数组都可以推广到高维。

Share