[LeetCode] 307. Range Sum Query - Mutable 解题思路

Given an integer array nums, find the sum of the elements between indices i and j (i ≤ j), inclusive.

The update(i, val) function modifies nums by updating the element at index i to val.

Example:

Given nums = [1, 3, 5]

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

Note:

1. The array is only modifiable by the update function.
2. You may assume the number of calls to update and sumRange function is distributed evenly.

问题：给定一个固定长度的数组，可以更新元素的值，求给定子数组的元素和。求和与更新操作交替进行。

解决方案

方案1，每次求和，直接遍历子数组进行求和。每次更新，直接根据下标更新元素值。求和操作时间复杂度为 O(n), 更新操作时间复杂度为O(1)

方案2，采用线段树存储原数组以及中间结果。

例子：输入数组{1, 3, 5, 7, 9, 11} 对应的线段树数据结构如下，叶子节点存储输入数组的元素，非叶子节点存储一个区域的和。

 

线段树是一个 full binary tree，可以用数组来存储。数组下标和线段树的节点之间的关系如下：

对于节点 i,

• 其左子节点下标为 i*2+1
• 其右节点下标为 i*2+2
• 其父亲节点下标为 (i-1)/2

基于数组的线段树表示例子如下：

 

求和思路：计算子数组[l, r] 的和时，对于给定节点 node 有：

• 如果节点 node 代表的范围在 [l, r] 之内，则返回节点 node 的值
• 如果节点 node 代表的范围完全不在 [l, r] 之内，则返回 0
• 其他情况，节点 node 代表的范围一部分在 [l, r]之内，一部分不在之内，则对于节点 node 的左右子节点分别应用该规则进行处理。

更新思路：根据给定的下标，更新下标对应元素在线段树的叶子节点，并更新从该叶子节点到根节点路径上的所有祖先节点。

 

构建线段树：根据输入数组，求得线段树需要的节点值。例如 {1, 3, 5, 7, 9, 11, 4, 12, 20, 16, 36}

对节点值进行反序处理，则得到基于数组结构的线段树。例如 {36, 16, 20, 12, 4, 11, 9, 7, 5, 3, 1}

 

代码实现如下：

#include <vector>
using namespace std;struct TreeNode{
    int val;
    pair<int, int> idxRange;

    TreeNode(int val, int lRange, int rRange){
        this->val = val;
        this->idxRange = make_pair(lRange, rRange);
    }
};

class NumArray {

    vector<int> nums;
    vector<TreeNode *> nodesVec;
    vector<TreeNode *> treeVec;

    /**
     * calculate the nodes of the segment tree based on the input values in nums
     */
    void calculateNodes(){
        for(int i=0; i < nums.size(); i++){
            TreeNode *tn = new TreeNode(nums[i], i, i);
            this->nodesVec.push_back(tn);
        }

        for(int i =0; i + 1 < nodesVec.size(); i+=2){
            int val = nodesVec[i]->val + nodesVec[i+1]->val;
            int l = min(nodesVec[i]->idxRange.first, nodesVec[i+1]->idxRange.first);
            int r = max(nodesVec[i]->idxRange.second, nodesVec[i+1]->idxRange.second);
            TreeNode *tn = new TreeNode(val, l, r);
            nodesVec.push_back(tn);
        }
    }

    /**
     * build segment tree base on Vector.
     * For node i,
     * the left child index: i * 2 + 1
     * the right child index: i * 2 + 2
     * the parent index: (i - 1)/2
     */
    void buildVectorBasedSegmentTree(){
        for(int i=nodesVec.size() - 1; i >= 0; i--){
            treeVec.push_back(nodesVec[i]);
        }
    }

public:
    NumArray(){}
    virtual ~NumArray(){}

    /**
     * initialization
     */
    NumArray(vector<int> nums){
        this->nums = nums;
        calculateNodes();
        buildVectorBasedSegmentTree();
    }


    /**
     * update the value of the node in the index i in treeVec
     */
    void update(int i, int val){
        int leafIdx = treeVec.size() - ( 1+ i );
        int diff = val - treeVec[leafIdx]->val;
        updateNodes(leafIdx, diff);
    }

    /**
     * update the value of nodes in interval tree(segment tree) with the diff recursively.
     */
    void updateNodes(int idx, int diff){
        treeVec[idx]->val += diff;
        if(idx > 0){
            idx = (idx -1) / 2;
            updateNodes(idx, diff);
        }
    }

    int sumRange(int i, int j){
        return getSum(0, i, j);
    }

    /**
     * check the node in i-th index in treeVec, to calculate the sum of leaf
     */
    int getSum(int nodeIdx, int rangeL, int rangeR){
        TreeNode* node = treeVec[nodeIdx];
        if (rangeL <= node->idxRange.first && node->idxRange.second <= rangeR){
            return node->val;
        }

        if (node->idxRange.second < rangeL || rangeR < node->idxRange.first){
            return 0;
        }
        int nodeIdxL = nodeIdx * 2 + 1;
        int nodeIdxR = nodeIdx * 2 + 2;
        return getSum(nodeIdxL, rangeL, rangeR) + getSum(nodeIdxR, rangeL, rangeR);
    }
};

 

 

Reference:

Segment Tree | Set 1 (Sum of given range), geeksforgeeks

Segment tree, wikipedia