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

#ifndef NUMARRAY_H_
#define NUMARRAY_H_
#include <vector>
using namespace std;

namespace segmentTree {

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);
    }

    /**
     *
     */
    string toString(){
        string res = "";
        res += to_string(nodesVec.size()) + " | ";
        for(vector<TreeNode*>::iterator it = nodesVec.begin(); it != nodesVec.end(); it++){
            TreeNode *node = *it;
            res += to_string(node->val) + "(" + to_string(node->idxRange.first) + "," + to_string(node->idxRange.second) + ")" + ", ";
        }
        res += "\n";

        res += to_string(treeVec.size()) + " | ";
        for(vector<TreeNode*>::iterator it = treeVec.begin(); it != treeVec.end(); it++){
            TreeNode *node = *it;
            res += to_string(node->val) + "(" + to_string(node->idxRange.first) + "," + to_string(node->idxRange.second) + ")" + ", ";
        }
        return res;
    }
};

} /* namespace segmentTree */

#endif /* NUMARRAY_H_ */