Interval Sum II 207

Question

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), 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.

Challenge

O(logN) time for query and modify.

Solution

水题，就是把query sum和modify结合起来而已。

代码如下：

public class Solution {
    /* you may need to use some attributes here */
    class SegmentTreeNode{
        int start;
        int end;
        long sum;
        SegmentTreeNode left;
        SegmentTreeNode right;
        public SegmentTreeNode(int start, int end, long sum){
            this.start = start;
            this.end = end;
            this.sum = sum;
            left = right = null;
        }
    }

    SegmentTreeNode root;

    /**
     * @param A: An integer array
     */
    public Solution(int[] A) {
        // write your code here
        if(A == null || A.length == 0){
            root = null;
        }else{
            root = SegmentTreeBuilder(A, 0, A.length - 1);
        }
    }

    private SegmentTreeNode SegmentTreeBuilder(int[] A, int start, int end){
        if(start == end){
            return new SegmentTreeNode(start, end, (long)A[start]);
        }

        int mid = (start + end) / 2;
        SegmentTreeNode root = new SegmentTreeNode(start, end, 0);
        SegmentTreeNode left = SegmentTreeBuilder(A, start, mid);
        SegmentTreeNode right = SegmentTreeBuilder(A, mid + 1, end);

        root.left = left;
        root.right = right;
        root.sum = left.sum + right.sum;

        return root;
    }

    /**
     * @param start, end: Indices
     * @return: The sum from start to end
     */
    public long query(int start, int end) {
        // write your code here
        if(root == null || start > end || start > root.end || end < root.start){
            return -1;
        }

        return queryHelper(root, start, end);
    }

    public long queryHelper(SegmentTreeNode r, int start, int end){
        if((start == r.start && end == r.end) || (r.start == r.end)){
            return r.sum;
        }

        int mid = (r.start + r.end) / 2;
        //跨越中点
        if(start <= mid && end >= mid + 1){
            long left = queryHelper(r.left, start, mid);
            long right = queryHelper(r.right, mid + 1, end);
            return left + right;
        }
        //左边
        if(end <= mid){
            return queryHelper(r.left, start, end);
        }
        //右边
        return queryHelper(r.right, start, end);
    }

    /**
     * @param index, value: modify A[index] to value.
     */
    public void modify(int index, int value) {
        // write your code here
        if(root == null || index < root.start || index > root.end){
            return;
        }

        modifyHelper(root, index, value);
    }

    public void modifyHelper(SegmentTreeNode root, int index, int value) {
        // write your code here
        if(root.start == index && root.end == index){
            root.sum = value;
            return;
        }

        int mid = (root.start + root.end) / 2;
        if(index <= mid){
            modifyHelper(root.left, index, value);
        }else{
            modifyHelper(root.right, index, value);
        }

        root.sum = root.left.sum + root.right.sum;
    }
}