Given an integer array in the construct method, implement two methods
query(start, end)
and modify(index, value)
:- For query(start, end), return the sum from index start to index end in the given array.
- For modify(index, value), modify the number in the given index to value
Example
Given array A =
[1,2,7,8,5]
.query(0, 2)
, return10
.modify(0, 4)
, change A[0] from 1 to 4.query(0, 1)
, return6
.modify(2, 1)
, change A[2] from 7 to 1.query(2, 4)
, return14
.
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); } };