Range Sum Query - Mutable

Posted 2020-08-11 Sheryl Wang

class SegmentTreeNode(object):
    def __init__(self, start, end):
        self.start = start
        self.end = end
        self.sum = 0
        self.left = None
        self.right = None
        
class NumArray(object):
    def __init__(self, nums):
        """
        initialize your data structure here.
        :type nums: List[int]
        """
        n = len(nums)
        self.root = self.buildTree(nums, 0, n-1)
        

    def update(self, i, val):
        """
        :type i: int
        :type val: int
        :rtype: int
        """
        self.modifyTree(i,val, self.root)
        

    def sumRange(self, i, j):
        """
        sum of elements nums[i..j], inclusive.
        :type i: int
        :type j: int
        :rtype: int
        """
        return self.queryTree(i, j, self.root)
        
    def buildTree(self, nums, start, end):
        if start > end: return None#not valid
        root = SegmentTreeNode(start, end)
        if start == end:
            root.sum = nums[start]
            return root
        mid = start + (end - start)/2
        root.left = self.buildTree(nums, start, mid)
        root.right = self.buildTree(nums, mid+1, end)
        root.sum = root.left.sum + root.right.sum
        return root
        
    def modifyTree(self, i, val, root):
        if not root:
            return 0
        if root.start==i and root.end == i: 
            diff = val - root.sum
            root.sum = val
            return diff
        mid = (root.start + root.end)/2
        if i > mid:
            diff = self.modifyTree(i, val, root.right)
        else:
            diff = self.modifyTree(i, val, root.left)
        root.sum = root.sum + diff
        return diff
        
    def queryTree(self, i, j, root):
        if not root:
            return 0
        if root.start == i and root.end == j:
            return root.sum
        mid = (root.start + root.end)/2
        if i > mid:
            return self.queryTree(i, j, root.right)
        if j <= mid:
            return self.queryTree(i, j, root.left)
            
        return self.queryTree(i, mid, root.left)+self.queryTree(mid+1, j , root.right)
        
        

# Your NumArray object will be instantiated and called as such:
# numArray = NumArray(nums)
# numArray.sumRange(0, 1)
# numArray.update(1, 10)
# numArray.sumRange(1, 2)