Monday, May 11, 2015

[LintCode] Interval Sum II

Given an integer array in the construct method, implement two methods query(start, end) and modify(index, value):
  • For query(startend), return the sum from index start to index end in the given array.
  • For modify(indexvalue), modify the number in the given index to value
Example
Given array A = [1,2,7,8,5].
  • query(0, 2), return 10.
  • modify(0, 4), change A[0] from 1 to 4.
  • query(0, 1), return 6.
  • modify(2, 1), change A[2] from 7 to 1.
  • query(2, 4), return 14.
class Solution 
{
    class SegmentTreeNode 
    {
    public:
        int start, end, sum;
        SegmentTreeNode *left, *right;
        SegmentTreeNode(int start, int end, int sum) {
            this->start = start;
            this->end = end;
            this->sum = sum;
            this->left = this->right = NULL;
        }
    };
    SegmentTreeNode *build(vector<int>&A, int start, int end)
    {
        if(start == end) return new SegmentTreeNode(start, end, A[start]);
        int mid = (start+end)/2;
        SegmentTreeNode *left = build(A, start, mid);
        SegmentTreeNode *right = build(A, mid+1, end);
        SegmentTreeNode *node = new SegmentTreeNode(start, end, left->sum+right->sum);
        node->left = left;
        node->right = right;
        return node;
    }
    SegmentTreeNode *root;
    long long query(SegmentTreeNode* root, int start, int end)
    {
        if(start == root->start && end == root->end) return root->sum;
        int mid = (root->start + root->end)/2;
        if(end <= mid) return query(root->left, start, end);
        if(start > mid) return query(root->right, start, end);
        return query(root->left, start, mid) + query(mid+1, end);
    }
    int modify(SegmentTreeNode *root, int index, int value) {
    
        if(root->start == root->end && index == root->start) 
        {
            int tmp = value - root->sum;
            root->sum = value;
            return tmp;
        }
        
        int mid = (root->start + root->end)/2;
        int diff;
        if(index <= mid) diff = modify(root->left, index, value);
        else diff = modify(root->right, index, value);
        root->sum += diff;
        return diff;
    }
public:
    Solution(vector<int> A) {
        if(A.size() > 0) root = build(A, 0, A.size()-1);
        else root = NULL;
    }
    long long query(int start, int end) {
        if(!root) return -1;
        return query(root, start, end);
    }
    void modify(int index, int value) {
        if(!root) return;
        modify(root, index, value);
    }
};