Segment Tree

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 ceil$
• Segment Tree的高度为$lceil log_2n ceil$, 因此该数组的大小为$2 * 2 ^{lceil log_2n ceil} - 1$

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

#include <math.h>
#include <vector>
#include <iostream>
using namespace std;

class segTree{
public:
    vector<int> tree;
    int n;
    segTree(vector<int>& arr){
        n = arr.size();
        int treeSize = 2 * pow(2, ceil(log2(double(n)))) - 1;
        tree.resize(treeSize);
        buildSegTree(arr, 0, 0, n - 1);
    }
    
    
    void buildSegTree(vector<int>& arr, int treeIndex, int lo, int hi){
        if(lo == hi){
            tree[treeIndex] = arr[lo];
            return;
        }
        int mid = lo + (hi - lo) / 2;
        buildSegTree(arr, 2 * treeIndex + 1, lo, mid);
        buildSegTree(arr, 2 * treeIndex + 2, mid + 1, hi);
        tree[treeIndex] = merge(tree[2 * treeIndex + 1], tree[2 * treeIndex + 2]);
    }


    int querySegTree(int treeIndex, int lo, int hi, int i, int j){
        if(lo > j || hi < i)
            return 0;
        if(i <= lo && j >= hi)
            return tree[treeIndex];
        
        int mid = lo + (hi - lo) / 2;
        
        if(i > mid)
            return querySegTree(2 * treeIndex + 2, mid + 1, hi, i, j);
        else if(j <= mid)
            return querySegTree(2 * treeIndex + 1, lo, mid, i, j);
        
        int leftQuery = querySegTree(2 * treeIndex + 1, lo, mid, i, mid);
        int rightQuery = querySegTree(2 * treeIndex + 2, mid + 1, hi, mid + 1, j);

        return merge(leftQuery, rightQuery);
    }


    void updateValSegTree(int treeIndex, int lo, int hi, int arrIndex, int val){
        if(lo == hi){
            tree[treeIndex] = val;
            return;
        }

        int mid = lo + (hi - lo) / 2;

        if(arrIndex > mid)
            updateValSegTree(2 * treeIndex + 2, mid + 1, hi, arrIndex, val);
        else if(arrIndex <= mid)
            updateValSegTree(2 * treeIndex + 1, lo, mid, arrIndex, val);
        
        tree[treeIndex] = merge(tree[2 * treeIndex + 1] , tree[2 * treeIndex + 2]);
    }


    int merge(int& v1, int& v2){
        return v1 + v2;
    }
};


int main(){
    vector<int> arr1 = {1, 3, 5, 7, 9, 11};
    segTree test(arr1);
    for(int item : test.tree)
        cout << item << " ";
    cout << endl;
    int sum = test.querySegTree(0, 0, test.n - 1, 0, 2);
    cout << "sum = " << sum << endl;
    test.updateValSegTree(0, 0, test.n - 1, 1, 4);
    int updatedSum  = test.querySegTree(0, 0, test.n - 1, 0, 2);
    cout << "updated sum = " << updatedSum << endl;
    return 0;
}