Binary Indexed Tree(树状数组) / Segment Tree (线段树)

Range Sum Query - Mutable (LeetCode)
多用于高效计算数列的前缀和， 区间和
在O(log n)时间内得到任意前缀和，并同时支持在O(log n)时间内动态单点值的修改

// Binary Indexed Tree 树状数组
// 86 ms
class NumArray {
public:
    NumArray(vector<int> nums) {
        array.resize(nums.size() + 1, 0); // Use array of length n + 1, ignoring index 0
        for (int i = 0; i < nums.size(); i++)
            add(i + 1, nums[i]);
    }

    void update(int i, int val) {
        int origin_val = sum(i + 1) - sum(i);
        add(i + 1, val - origin_val);
    }
    
    int sumRange(int i, int j) {
        return sum(j + 1) - sum(i);
    }
private:
    vector<int> array;
    void add(int i, int val) {
        while (i < array.size()) {
            array[i] += val;
            i += (i & -i); // Extract last set bit: x & (-x)
        }
    }
    int sum(int i) {
        int res = 0;
        while (i) {
            res += array[i];
            i -= (i & -i);
        }
        return res;
    }
};

// Segment Tree 线段树
// 133 ms
class NumArray {
public:
    NumArray(vector<int> nums) {
        n = nums.size();
        if (n > 0) {
            // Height of segment tree
            int x = (int)(ceil(log2(n))); 
            // Maximum size of segment tree
            int max_size = 2 * (int)pow(2, x) - 1; 
            st.resize(max_size); // Important!: st.resize(2 * n - 1); 是错误的 因为线段树不一定是完全二叉树
            initialize(0, 0, n - 1, nums);
        }
    }
    
    void update(int i, int val) {
        int origin_val = sumRange(i, i);
        add(i, val - origin_val);
    }
    
    int sumRange(int i, int j) {
        return sumRange(0, i, j, 0, n - 1);
    }
private:
    int n;
    vector<int> st;
    int initialize(int idx, int l, int r, vector<int> &nums) {
        if (l == r) {
            st[idx] = nums[l];
        } else {
            int mid = (l + r) / 2;
            st[idx] = initialize(2 * idx + 1, l, mid, nums) + initialize(2 * idx + 2, mid + 1, r, nums);
        }
        return st[idx];
    }
    void add(int i, int val) {
        int l = 0, r = n - 1, idx = 0;
        while (l < r) {
            st[idx] += val;
            int mid = (l + r) / 2;
            if (i <= mid) {
                idx = 2 * idx + 1;
                r = mid;
            } else {
                idx = 2 * idx + 2;
                l = mid + 1;
            }
        }
        st[idx] += val;
    }
    int sumRange(int idx, int i, int j, int l, int r) {
        if (l > r || j < l || r < i)
            return 0;
        if (i <= l && r <= j)
            return st[idx];
        int mid = (l + r) / 2;
        return sumRange(idx * 2 + 1, i, j, l, mid) + sumRange(idx * 2 + 2, i, j, mid + 1, r);
    }
};

简书原创，转载请联系作者