Interval Sum II

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.

Thoughts

Solution

memory limit exceeded ...

class Solution: 

    # @param A: An integer list
    def __init__(self, A):
        # write your code here
        self.segT = SegmentTree()
        length = len(A)
        self.root = self.segT.build(0, length)
        for idx, num in enumerate(A):
            self.segT.modify(self.root, idx, num)

# @param start, end: Indices
    # @return: The sum from start to end
    def query(self, start, end):
        # write your code here
        res = 0
        res = self.segT.query(self.root, start, end)
        return res

    # @param index, value: modify A[index] to value.
    def modify(self, index, value):
        # write your code here
        self.segT.modify(self.root, index, value)

class SegmentTree:

    def __init__(self):
        self.root = None

    def build(self, start, end):
        if start > end:
            return None
        root = SegmentTreeNode(start, end, 0)

        if start < end:
            mid = (start + end) / 2
            root.left = self.build(start, mid)
            root.right = self.build(mid+1, end)
        else:
            root.count = 0

        return root

    def query(self, root, start, end):
        if not root:
            return 0
        rStart = root.start
        rEnd = root.end

        if (rStart > end) or ( rEnd < start):
            return 0

        #should be >= and <=
        #if query range if larger than the segment tree range, just return the count
        if rStart >= start and rEnd <= end:
            return root.count

        leftCount = 0
        rightCount = 0

        mid = (rStart + rEnd) / 2
        if start <= mid:
            if end > mid:
                #some range are in right part
                leftCount = self.query(root.left, start, mid)
            else:
                leftCount = self.query(root.left, start, end)
        if mid < end:
            if start <= mid:
                #some range are in left part
                #right part should start with mid + 1
                rightCount = self.query(root.right, mid + 1, end)
            else:
                rightCount = self.query(root.right, start, end)
        return leftCount + rightCount

    def modify(self, root, index, value):
        if root.start == index and root.end == index:
            root.count = value
            return

        mid = (root.start + root.end) / 2
        if root.start <= index and index <= mid:
            self.modify(root.left, index, value)
        elif root.end >= index and index > mid:
            self.modify(root.right, index, value)

        root.count = root.left.count + root.right.count

java code

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

     class SegmentTreeNode {
        public int start, end;
        public int sum;
        public SegmentTreeNode left, right;
        public SegmentTreeNode(int start, int end, int sum) {
              this.start = start;
              this.end = end;
              this.sum = sum;
              this.left = this.right = null;
        }
    }
    SegmentTreeNode root;
    public SegmentTreeNode build(int start, int end, int[] A) {
        // write your code here
        if(start > end) {  // check core case
            return null;
        }

        SegmentTreeNode root = new SegmentTreeNode(start, end, 0);

        if(start != end) {
            int mid = (start + end) / 2;
            root.left = build(start, mid, A);
            root.right = build(mid+1, end, A);

            root.sum = root.left.sum + root.right.sum;
        } else {
            root.sum =  A[start];

        }
        return root;
    }
    public int querySegmentTree(SegmentTreeNode root, int start, int end) {
        // write your code here
        if(start == root.start && root.end == end) { // 相等 
            return root.sum;
        }


        int mid = (root.start + root.end)/2;
        int leftsum = 0, rightsum = 0;
        // 左子区
        if(start <= mid) {
            if( mid < end) { // 分裂 
                leftsum =  querySegmentTree(root.left, start, mid);
            } else { // 包含 
                leftsum = querySegmentTree(root.left, start, end);
            }
        }
        // 右子区
        if(mid < end) { // 分裂 3
            if(start <= mid) {
                rightsum = querySegmentTree(root.right, mid+1, end);
            } else { //  包含 
                rightsum = querySegmentTree(root.right, start, end);
            } 
        }  
        // else 就是不相交
        return leftsum + rightsum;
    }
    public void modifySegmentTree(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(root.start <= index && index <=mid) {
            modifySegmentTree(root.left, index, value);
        }

        if(mid < index && index <= root.end) {
            modifySegmentTree(root.right, index, value);
        }
        //更新
        root.sum = root.left.sum + root.right.sum;
    }
    /**
     * @param A: An integer array
     */
    public Solution(int[] A) {
        // write your code here
        root = build(0, A.length-1, A);
    }

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

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