线段树求区间和

视频链接:https://www.bilibili.com/video/BV1cb411t7AM/?spm_id_from=333.788.top_right_bar_window_history.content.click&vd_source=df505e255a0c2e428b302beccab437c3

测试代码:

  1 #include <string>
  2 #include <unordered_set>
  3 #include <iostream>
  4 #include <algorithm>
  5 
  6 int length;
  7 std::vector<int> arr;
  8 std::vector<int> tree;
  9 
 10 // start -> end = sum 
 11 void buildTree(int node, int start, int end) {
 12     if (start == end) {
 13         tree[node] = arr[start];
 14     }
 15     else {
 16         // 计算左节点和右节点
 17         int leftNode = 2 * node + 1;
 18         int rightNode = 2 * node + 2;
 19         // 计算中间树节点
 20         int mid = (start + end) / 2;
 21         // 计算当前节点左节点的值
 22         buildTree(leftNode, start, mid);
 23         // 计算当前节点右节点的值
 24         buildTree(rightNode, mid + 1, end);
 25         tree[node] = tree[leftNode] + tree[rightNode];
 26     }
 27 }
 28 
 29 void updateTree(int node, int start, int end, int idx, int val) {
 30     if (start == end) {
 31         arr[idx] = val;
 32         tree[node] = val;
 33     }
 34     else {
 35         int mid = (start + end) / 2;
 36         int leftNode = 2 * node + 1;
 37         int rightNode = 2 * node + 2;
 38         if (idx >= start && idx <= mid) {
 39             updateTree(leftNode, start, mid, idx, val);
 40         }
 41         else {
 42             updateTree(rightNode, mid + 1, end, idx, val);
 43         }
 44         tree[node] = tree[leftNode] + tree[rightNode];
 45     }
 46 
 47 }
 48 
 49 int queryTree(int node, int start, int end, int left, int right) {
 50     if (right < start || end < left) {
 51         return 0;
 52     }
 53     else if (left <= start && end <= right) {
 54         return tree[node];
 55     }
 56     else if (start == end) {
 57         return tree[node];
 58     }
 59     else {
 60         int mid = (start + end) / 2;
 61         int leftNode = 2 * node + 1;
 62         int rightNode = 2 * node + 2;
 63         int sumLeft = queryTree(leftNode, start, mid, left, right);
 64         int sumRight = queryTree(rightNode, mid + 1, end, left, right);
 65         return sumLeft + sumRight;
 66     }
 67 
 68 }
 69 
 70 void NumArray(std::vector<int>& nums) {
 71     length = nums.size();
 72     arr.assign(nums.begin(), nums.end());
 73     // 通过计算,tree数组需要的长度小于4 * length
 74     tree.assign(length << 2, 0);
 75 
 76     buildTree(0, 0, length - 1);
 77 }
 78 
 79 void update(int index, int val) {
 80     updateTree(0, 0, length - 1, index, val);
 81 }
 82 
 83 int sumRange(int left, int right) {
 84     return queryTree(0, 0, length - 1, left, right);
 85 }
 86 
 87 int main() {
 88     std::vector<int> nums = { 1, 3, 5, 7, 9, 11 };
 89     NumArray(nums);
 90     for (auto st : tree) {
 91         std::cout << st << " ";
 92     }
 93     std::cout << std::endl;
 94 
 95     updateTree(0, 0, 5, 4, 6);
 96     for (auto st : tree) {
 97         std::cout << st << " ";
 98     }
 99     std::cout << std::endl;
100 
101     int s = queryTree(0, 0, 5, 2, 5);
102     std::cout << "2 - 5 sum: " << s << std::endl;
103 
104     return 0;    
105 }

 

posted @ 2023-02-28 16:41  风影旋新月  阅读(13)  评论(0编辑  收藏  举报