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