Lintcode: Interval Sum II

非常有借鉴意义的interval tree。

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

    public Solution(int[] A) {
        root = constructTree(A, 0, A.length - 1);
    }
			    
    private SumTreeNode constructTree(int[] a, int start, int end) {
    	if (start > end) {
    		return null;
    	}
    	if (start == end) {
    		return new SumTreeNode(start, start, (long)a[start]);
    	}
    	int mid = (start + end) / 2;
    	SumTreeNode root = new SumTreeNode(start, end);
    	root.left = constructTree(a, start, mid);
    	root.right = constructTree(a, mid + 1, end);
    	long sum = 0;
    	if (root.left != null) {
    		sum += root.left.sum;
    	}
    	if (root.right != null) {
    		sum += root.right.sum;
    	}
    	root.sum = sum;
    	return root;
    }
    /**
     * @param start, end: Indices
     * @return: The sum from start to end
     */
    public long query(int start, int end) {
        return subQuery(root, start, end);
    }
    
    public long subQuery(SumTreeNode root, int start, int end) {
    	if (root == null || root.end < start || root.start > end) {
        	return 0;
        }
        if (root.start >= start && root.end <= end) {
        	return root.sum;
        }
        return subQuery(root.left, start, end) + subQuery(root.right, start, end);
    }
    /**
     * @param index, value: modify A[index] to value.
     */
    public void modify(int index, int value) {
        subModify(root, index, value);
    }
    
    private void subModify(SumTreeNode root, int index, int value) {
    	if (root == null || root.end < index || root.start > index) {
        	return;
        }
    	if (root.start == root.end) {
    		root.sum = value;
    		return;
    	}
    	if ((root.start + root.end) / 2 >= index) {
    		subModify(root.left, index, value);
    	} else {
    		subModify(root.right, index, value);
    	}
    	long sum = 0;
    	if (root.left != null) {
    		sum += root.left.sum;
    	}
    	if (root.right != null) {
    		sum += root.right.sum;
    	}
    	root.sum = sum;
    }
}

Leave a comment