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:

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.

Solution:

This problem can be easily solved by Segment Tree.

Code:

public class Solution {
    class SegmentNode {
        int start;
        int end;
        long sum;
        SegmentNode left;
        SegmentNode right;
        public SegmentNode(int start, int end) {
            this.start = start;
            this.end = end;
            this.sum = 0;
        }
    }

    private SegmentNode root;

    public Solution(int[] A) {
        if (A == null || A.length == 0) {
            root = null;
        } else {
            root = buildTree(A);
            for (int i = 0; i < A.length; i++) {
                modify(i, A[i]);
            }
        }

    }

    private SegmentNode buildTree(int[] A) {
        Stack<SegmentNode> stack = new Stack<>();
        SegmentNode root = new SegmentNode(0, A.length - 1);
        stack.push(root);
        while (!stack.empty()) {
            SegmentNode node = stack.pop();
            if (node.start == node.end) {
                continue;
            }
            int mid = node.start + (node.end - node.start) / 2;
            if (node.start <= mid) {
                node.left = new SegmentNode(node.start, mid);
                stack.push(node.left);
            } 
            if (mid + 1 <= node.end) {
                node.right = new SegmentNode(mid + 1, node.end);
                stack.push(node.right);
            }
        }
        return root;
    }

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

    private long queryHelper(SegmentNode node, int start, int end) {
        if (node == null || start > end) {
            return 0L;
        }
        if (start <= node.start && node.end <= end) {
            return node.sum;
        }
        long sum = 0;
        int mid = node.start + (node.end - node.start) / 2;
        if (start <= mid) {
            if (mid < end) {
                sum += queryHelper(node.left, start, mid);
            } else {
                sum += queryHelper(node.left, start, end);
            }
        }
        if (mid < end) {
            if (start <= mid) {
                sum += queryHelper(node.right, mid + 1, end);
            } else {
                sum += queryHelper(node.right, start, end);
            }
        }
        return sum;
    }

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

    private void modifyHelper(SegmentNode node, int index, int value) {
        if (node == null) {
            return;
        }
        if (node.start == index && node.end == index) {
            node.sum = value;
            return;
        }
        int mid = node.start + (node.end - node.start) / 2;
        if (index <= mid) {
            modifyHelper(node.left, index, value);
        } else {
            modifyHelper(node.right, index, value);
        }
        if (node.left != null) {
            node.sum = node.left.sum;
        }
        if (node.right != null) {
            node.sum += node.right.sum;
        }
    }
}