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
.
分析:按题目提示,最好先做过线段树的构造,线段树的查询,线段树的修改三道题,因为这道题分别涉及线段树的构造/修改/查询。第一个一次通过的困难题。
public class Solution {
/* you may need to use some attributes here */
class SegmentTreeNode {
public int start, end;
long val;
public SegmentTreeNode left, right;
public SegmentTreeNode(int start, int end) {
this.start = start;
this.end = end;
this.left = this.right = null;
}
}
private SegmentTreeNode tree = null;
/**
* @param A: An integer array
*/
public Solution(int[] A) {
if(A == null) return;
tree = build(A, 0, A.length - 1);
}
private SegmentTreeNode build(int[] A, int start, int end) {
if(start > end) return null;
SegmentTreeNode res = new SegmentTreeNode(start, end);
if(start == end) {
res.val = A[start];
} else {
SegmentTreeNode left = build(A, start, (start + end) / 2);
SegmentTreeNode right = build(A, (start + end) / 2 + 1, end);
res.left = left;
res.right = right;
res.val = left.val + right.val;
}
return res;
}
/**
* @param start, end: Indices
* @return: The sum from start to end
*/
public long query(int start, int end) {
return query(tree, start, end);
}
private long query(SegmentTreeNode curr, int start, int end) {
if(curr.start == start && curr.end == end) return curr.val;
if(start > curr.end || end < curr.start) return 0;
long left = query(curr.left, start, Math.min(end, curr.left.end));
long right = query(curr.right, Math.max(start, curr.right.start), end);
return left + right;
}
/**
* @param index, value: modify A[index] to value.
*/
public void modify(int index, int value) {
modify(tree, index, value);
}
private long modify(SegmentTreeNode curr, int index, int value) {
if(index < curr.start || index > curr.end) return 0;
if(curr.start == curr.end && curr.start == index) {
long diff = value - curr.val;
curr.val = value;
return diff;
} else {
if((curr.start + curr.end) / 2 >= index) {
long diff = modify(curr.left, index, value);
curr.val = curr.val + diff;
return diff;
} else {
long diff = modify(curr.right, index, value);
curr.val = curr.val + diff;
return diff;
}
}
}
}