Interval Sum II

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

Notice

We suggest you finish problem Segment Tree Build, Segment Tree Query and Segment Tree Modify first.

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.

class SegmentTreeNode {  
    public int start, end, sum;
    public SegmentTreeNode left, right;

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

public class Solution {  
    SegmentTreeNode root;

    /**
     * @param A: An integer array
     */
    public Solution(int[] A) {
        if (A == null || A.length == 0) {
            return;
        }

        root = buildTree(0, A.length - 1, A);
    }

    SegmentTreeNode buildTree(int start, int end, int[] A) {
        if (start > end) {
            return null;
        }

        SegmentTreeNode root = new SegmentTreeNode(start, end);

        if (start == end) {
            root.sum = A[start];
        } else {
            int middle = start + (end - start) / 2;

            root.left = buildTree(start, middle, A);
            root.right = buildTree(middle + 1, end, A);

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

        return root;
    }

    /**
     * @param start, end: Indices
     * @return: The sum from start to end
     */
    public long query(int start, int end) {
        return query(root, start, end);
    }

    long query(SegmentTreeNode node, int start, int end) {
        if (node == null) {
            return 0;
        }

        if (start <= node.start && end >= node.end) {
            return node.sum;
        }

        int middle = node.start + (node.end - node.start) / 2;

        if (end <= middle) {
            return query(node.left, start, end);
        } else if (start > middle) {
            return query(node.right, start, end);
        } else {
            return query(node.left, start, middle) + 
                query(node.right, middle + 1, end);
        }
    }

    /**
     * @param index, value: modify A[index] to value.
     */
    public void modify(int index, int value) {
        modify(root, index, value);    
    }

    void modify(SegmentTreeNode root, int index, int value) {
        if (root == null) {
            return;
        }
        if (index < root.start || index > root.end) {
            return;
        }

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

        int middle = root.start + (root.end - root.start) / 2;

        if (index <= middle) {
            modify(root.left, index, value);
        } else {
            modify(root.right, index, value);
        }

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

Hope this helps,
Michael