Segment Tree

Posted 2022-08-31 betaa

Leetcode上一道题，给定一个整数数组，要实现：

• 求[i, j]所有元素的和，0 <= i <= j <= n - 1，sumRange(i, j)
• 数组的元素会被修改， update(i, val)
• 这两个函数会被均匀的调用很多次

最简单的方法是求和O(n)，修改元素O(1)，时间复杂度太大，使用Segment Tree可以将二者的时间复杂度均变为O(logn)

Segment Tree:

• 叶子节点是输入数组中的所有元素
• 内部节点是其孩子节点所带信息的merge
• Segment Tree可以由数组实现，数组索引i的左孩子为$2 * i + 1$，右孩子为$2 * i + 2$，父节点在$\lceil (i - 1) / 2 \rceil$
• Segment Tree的高度为$\lceil log_2n \rceil$, 因此该数组的大小为$2 * 2 ^\lceil log_2n\rceil - 1$

一下代码是上述题目的C++递归实现：

 1 #include <math.h>
 2 #include <vector>
 3 #include <iostream>
 4 using namespace std;
 5 
 6 class segTree
 7 public:
 8     vector<int> tree;
 9     int n;
10     segTree(vector<int>& arr)
11         n = arr.size();
12         int treeSize = 2 * pow(2, ceil(log2(double(n)))) - 1;
13         tree.resize(treeSize);
14         buildSegTree(arr, 0, 0, n - 1);
15     
16     
17     
18     void buildSegTree(vector<int>& arr, int treeIndex, int lo, int hi)
19         if(lo == hi)
20             tree[treeIndex] = arr[lo];
21             return;
22         
23         int mid = lo + (hi - lo) / 2;
24         buildSegTree(arr, 2 * treeIndex + 1, lo, mid);
25         buildSegTree(arr, 2 * treeIndex + 2, mid + 1, hi);
26         tree[treeIndex] = merge(tree[2 * treeIndex + 1], tree[2 * treeIndex + 2]);
27     
28 
29 
30     int querySegTree(int treeIndex, int lo, int hi, int i, int j)
31         if(lo > j || hi < i)
32             return 0;
33         if(i <= lo && j >= hi)
34             return tree[treeIndex];
35         
36         int mid = lo + (hi - lo) / 2;
37         
38         if(i > mid)
39             return querySegTree(2 * treeIndex + 2, mid + 1, hi, i, j);
40         else if(j <= mid)
41             return querySegTree(2 * treeIndex + 1, lo, mid, i, j);
42         
43         int leftQuery = querySegTree(2 * treeIndex + 1, lo, mid, i, mid);
44         int rightQuery = querySegTree(2 * treeIndex + 2, mid + 1, hi, mid + 1, j);
45 
46         return merge(leftQuery, rightQuery);
47     
48 
49 
50     void updateValSegTree(int treeIndex, int lo, int hi, int arrIndex, int val)
51         if(lo == hi)
52             tree[treeIndex] = val;
53             return;
54         
55 
56         int mid = lo + (hi - lo) / 2;
57 
58         if(arrIndex > mid)
59             updateValSegTree(2 * treeIndex + 2, mid + 1, hi, arrIndex, val);
60         else if(arrIndex <= mid)
61             updateValSegTree(2 * treeIndex + 1, lo, mid, arrIndex, val);
62         
63         tree[treeIndex] = merge(tree[2 * treeIndex + 1] , tree[2 * treeIndex + 2]);
64     
65 
66 
67     int merge(int& v1, int& v2)
68         return v1 + v2;
69     
70 ;
71 
72 
73 int main()
74     vector<int> arr1 = 1, 3, 5, 7, 9, 11;
75     segTree test(arr1);
76     for(int item : test.tree)
77         cout << item << " ";
78     cout << endl;
79     int sum = test.querySegTree(0, 0, test.n - 1, 0, 2);
80     cout << "sum = " << sum << endl;
81     test.updateValSegTree(0, 0, test.n - 1, 1, 4);
82     int updatedSum  = test.querySegTree(0, 0, test.n - 1, 0, 2);
83     cout << "updated sum = " << updatedSum << endl;
84     return 0;
85