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

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

Thought

First, we need to build a segment tree, so it is necessary to declare a class for the node of the segment tree. Then, use recursion for query function. The modify function is almost similar to Segment Tree Modify.

• Build the segment tree: build the root node with the node class and use recursion to build the left and right child. Don't forget to update the value in the node.
• Query: Use recursion and compare the mid index of the segment node with the input start and end variable.
• Modify: Use a stack to store all the nodes connected the root and the desired node. Update the values of these nodes from bottom to the top.

Code

class SegmentNode:
    def __init__(self, value, start, end):
        self.value = value
        self.start = start
        self.end = end
        self.left = None
        self.right = None



class Solution:    

    # @param A: An integer list
    def __init__(self, A):
        # write your code here
        self.segmentTree = self.buildSegmentTree(A, 0, len(A) - 1)

    def buildSegmentTree(self, nums, start, end):
        if start > end:
            return None
        if start == end:
            return SegmentNode(nums[start], start, end)
        root = SegmentNode(0, start, end)
        mid = (start + end) / 2
        left = self.buildSegmentTree(nums, start, mid)
        right = self.buildSegmentTree(nums, mid + 1, end)
        root.left, root.right = left, right
        root.value = left.value + right.value
        return root



    # @param start, end: Indices
    # @return: The sum from start to end
    def query(self, start, end):
        # write your code here
        def queryHelper(node, start, end):
            if node.start ==start and node.end == end:
                return node.value
            mid = (node.start + node.end) / 2
            leftSum, rightSum = 0, 0
            if start <= mid:
                if end > mid:
                    leftSum = queryHelper(node.left, start, mid)
                else:
                    leftSum = queryHelper(node.left, start, end)
            if end > mid:
                if start <= mid:
                    rightSum = queryHelper(node.right, mid + 1, end)
                else:
                    rightSum = queryHelper(node.right, start, end)
            return leftSum + rightSum

        return queryHelper(self.segmentTree, start, end)

    # @param index, value: modify A[index] to value.
    def modify(self, index, value):
        # write your code here
        if self.segmentTree is None:
            return
        ptr = self.segmentTree
        stack = [ptr]
        while ptr.start < ptr.end:
            if ptr.left.end >= index:
                ptr = ptr.left
            else:
                ptr = ptr.right
            stack.append(ptr)
        while stack:
            ptr = stack.pop()
            if ptr.start == index and ptr.end == index:
                ptr.value = value
            else:
                ptr.value = ptr.left.value + ptr.right.value