LeetCode 307. 区域和检索 - 数组可修改

题目描述

给定一个整数数组  nums，求出数组从索引 i 到 j  (i ≤ j) 范围内元素的总和，包含 i,  j 两点。

update(i, val) 函数可以通过将下标为 i 的数值更新为 val，从而对数列进行修改。

示例:

Given nums = [1, 3, 5]

sumRange(0, 2) -> 9
update(1, 2)
sumRange(0, 2) -> 8

说明:

1. 数组仅可以在 update 函数下进行修改。
2. 你可以假设 update 函数与 sumRange 函数的调用次数是均匀分布的。

问题分析

据说此题是线段树的典型题，所以用线段树来做。由于线段树是满二叉树，所以其实际有效节点数是给定数组大小n的2倍减1，即2n-1，但是创建表示线段树的数组stree时，要创建成大小是2n的，第一个元素没有用，这样用节点的下标 /2 就可以得到其父节点的下标了，很方便。比如数组 [1, 3, 5, 7] 表示成线段树如下：对应的线段树数组为：[0, 16, 4, 12, 1, 3, 5, 7]。

           [0, 3]
             16
       /            \
    [0, 1]        [2, 3]
      4             12
   /     \       /     \
[0, 0] [1, 1] [2, 2] [3, 3]
   1      3      5      7

 所以本题我们在构造函数中创建出线段树stree，先在数组的后半部分创建叶子节点，然后在数组的前半部分创建其他节点。修改数组某元素的值时，我们先把其对应的叶子节点的值修改了，然后由于树中受影响的节点是从这个叶子节点往上的所有父节点，所以接下来我们只需更新这些父节点即可。更新时 i / 2 得到 i 的父节点，i ^ 1 得到 i 的兄弟节点，所以更新式为：stree[i / 2] = stree[i] + stree[i ^ 1]。当要返回 i 到 j 的范围内的元素和时，如果 i 本身是右孩子，那么将sum加上它本身，然后将i++；如果 j 本身是左孩子，那么将sum加上它本身，然后将j--。然后将 i 和 j 都除以2指到它们的父节点上。在 i 小于等于 j 时进行循环，当循环退出时，sum就是 i 到 j 范围内的元素和，返回即可。

代码实现

class NumArray {
private:
    int n;
    vector<int> stree;
    
public:
    NumArray(vector<int>& nums) {
        n = nums.size();
        stree.resize(2 * n);
        for(int i = n; i < 2 * n; i++)
            stree[i] = nums[i - n];
        for(int i = n - 1; i > 0; i--) 
            stree[i] = stree[i * 2] + stree[i * 2 + 1];
    }
    
    void update(int i, int val) {
        i = i + n;
        stree[i] = val;
        while(i > 0){
            stree[i / 2] = stree[i] + stree[i ^ 1];
            i = i / 2;
        }
    }
    
    int sumRange(int i, int j) {
        int sum = 0;
        i = i + n;
        j = j + n;
        while(i <= j){
            if((i & 1) == 1)
                sum += stree[i++];
            if((j & 1) == 0)
                sum += stree[j--];
            i = i / 2;
            j = j / 2;
        }
        return sum;
    }
};

/**
 * Your NumArray object will be instantiated and called as such:
 * NumArray* obj = new NumArray(nums);
 * obj->update(i,val);
 * int param_2 = obj->sumRange(i,j);
 */