Lintcode 207. Interval Sum II

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

Example 1

Input:
[1,2,7,8,5]
[query(0,2),modify(0,4),query(0,1),modify(2,1),query(2,4)]
Output: [10,6,14]
Explanation:
Given array A = [1,2,7,8,5].
After query(0, 2), 1 + 2 + 7 = 10,
After modify(0, 4), change A[0] from 1 to 4, A = [4,2,7,8,5].
After query(0, 1), 4 + 2 = 6.
After modify(2, 1), change A[2] from 7 to 1，A = [4,2,1,8,5].
After query(2, 4), 1 + 8 + 5 = 14.

Example 2

Input:
[1,2,3,4,5]
[query(0,0),query(1,2),quert(3,4)]
Output: [1,5,9]
Explantion:
1 = 1
2 + 3 = 5
4 + 5 = 9

Challenge

O(logN) time for query and modify.

Code (Java):

public class Solution {
    /* you may need to use some attributes here */
    SegmentTree segTree;

    /*
    * @param A: An integer array
    */
    public Solution(int[] A) {
        // do intialization if necessary
        segTree = new SegmentTree(A);
    }

    /*
     * @param start: An integer
     * @param end: An integer
     * @return: The sum from start to end
     */
    public long query(int start, int end) {
        return segTree.query(start, end);
    }

    /*
     * @param index: An integer
     * @param value: An integer
     * @return: nothing
     */
    public void modify(int index, int value) {
        // write your code here
        segTree.modify(index, value);
    }
}

class SegmentTree {
    private SegmentTreeNode root;

    public SegmentTree(int[] A) {
        root = buildSegmentTree(A);
    }

    public long query(int start, int end) {
        return queryHelper(root, start, end);
    }

    public void modify(int index, int value) {
        modifyHelper(root, index, value);
    }

    private void modifyHelper(SegmentTreeNode root, int index, int value) {
        if (root.left == null && root.right == null && root.start == index && root.end == index) {
            root.sum = value;
            return;
        }

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

    private long queryHelper(SegmentTreeNode root, int start, int end) {
        if (root == null || start > end) {
            return 0;
        }

        if (start == root.start && end == root.end) {
            return root.sum;
        }

        int mid = root.start + (root.end - root.start) / 2;
        long leftSum = queryHelper(root.left, Math.max(root.start, start), Math.min(mid, end));
        long rightSum = queryHelper(root.right, Math.max(mid + 1, start), Math.min(root.end, end));
        return leftSum + rightSum;
    }

    private SegmentTreeNode buildSegmentTree(int[] A) {
        if (A == null || A.length == 0) {
            return null;
        }
    
        return buildSegmentTreeHelper(A, 0, A.length - 1);
    }
    
    private SegmentTreeNode buildSegmentTreeHelper(int[] A, int start, int end) {
        if (start == end) {
            return new SegmentTreeNode(A[start], start, start);
        }
    
        int mid = start + (end - start) / 2;
        SegmentTreeNode leftChild = buildSegmentTreeHelper(A, start, mid);
        SegmentTreeNode rightChild = buildSegmentTreeHelper(A, mid + 1, end);
    
        SegmentTreeNode root = new SegmentTreeNode(leftChild.sum + rightChild.sum, start, end);
        root.left = leftChild;
        root.right = rightChild;
    
        return root;
    }

    class SegmentTreeNode {
        long sum;
        int start, end;
        SegmentTreeNode left, right;
        public SegmentTreeNode(long sum, int start, int end) {
            this.sum = sum;
            this.start = start;
            this.end = end;
            left = null;
            right = null;
        }
    }
}