Given an integer array in the construct method, implement two methodsquery(start, end)
andmodify(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
Notice
We suggest you finish problemSegment Tree Build,Segment Tree QueryandSegment Tree Modifyfirst.
Have you met this question in a real interview?
Yes
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
.
O(logN) time forquery
andmodify
.
class Solution {
public:
/\* you may need to use some attributes here \*/
/\*\*
\* @param A: An integer vector
\*/
class SegmentTreeNode {
public:
int start, end;
long long sum;
SegmentTreeNode \*left, \*right;
SegmentTreeNode\(int start, int end, int sum\) :start\(start\), end\(end\), sum\(sum\) {
left = NULL;
right = NULL;
}
};
SegmentTreeNode\* root;
Solution\(vector<int> A\) {
// write your code here
root = buildSegmentTree\(A, 0, A.size\(\) - 1\);
}
/\*\*
\* @param start, end: Indices
\* @return: The sum from start to end
\*/
long long query\(int start, int end\) {
// write your code here
return query\(root, start, end\);
}
long long query\(SegmentTreeNode\* root, int start, int end\) {
if \(start > end \|\| root == NULL\) {
return 0;
}
if \(start <= root->start && root->end <= end\) {
return root->sum;
}
long long leftSum = query\(root->left, start, end\);
long long rightSum = query\(root->right, start, end\);
return leftSum + rightSum;
}
/\*\*
\* @param index, value: modify A\[index\] to value.
\*/
void modify\(int index, int value\) {
// write your code here
modify\(root, index, value\);
}
void modify\(SegmentTreeNode\* root, int index, int value\) {
if \(root->start == root->end && index == root->start\) {
root->sum = value;
return;
}
if \(root->start <= index && index <= root->end\) {
modify\(root->left, index, value\);
modify\(root->right, index, value\);
root->sum = root->left->sum + root->right->sum;
}
}
SegmentTreeNode\* buildSegmentTree\(vector<int>& A, int start, int end\) {
if \(start > end\) {
return NULL;
}
if \(start == end\) {
return new SegmentTreeNode\(start, end, A\[start\]\);
}
int mid = start + \(end - start >> 1\);
SegmentTreeNode\* root = new SegmentTreeNode\(start, end, 0\);
root->left = buildSegmentTree\(A, start, mid\);
root->right = buildSegmentTree\(A, mid + 1, end\);
root->sum += root->left->sum;
root->sum += root->right->sum;
return root;
}
};