LeetCode #307: Range Sum Query - Mutable

Description:

Given an integer array nums, find the sum of the elements between indices i and j (i ≤ j), inclusive.

The update(i, val) function modifies nums by updating the element at index i to val.

Example:

Given nums = [1, 3, 5]

sumRange(0, 2) -> 9
update(1, 2)
sumRange(0, 2) -> 8

Note:

1. The array is only modifiable by the update function.
2. You may assume the number of calls to update and sumRange function is distributed evenly.

Solution:

• Use segmental tree to restore the sums between any indices
• Time complexity: update O(log(N)); sum O(log(N))
• Space complexity: O(N)

class NumArray:

    def __init__(self, nums: List[int]):
        #import numpy as np
        import math

        self.array = nums
        if len(nums) == 0:
            self.tree = []
        else:
            nlayer = int(math.ceil(math.log(len(nums), 2))) + 1
            self.tree = [None for i in  range(2**nlayer)]
            self.build_segment_tree(self.array, 0, len(self.array)-1, self.tree, 0)        

    def build_segment_tree(self, array, start, end, tree, inode):
        if len(array) > 0:
            if start == end:
                tree[inode] = array[start]
            else:
                #print(inode, start, end)
                mid = int((start+end)/2)
                self.build_segment_tree(array, start, mid, tree, inode*2+1)
                self.build_segment_tree(array, mid+1, end, tree, inode*2+2)
                tree[inode] = tree[inode*2+1] + tree[inode*2+2]
            return

    def update(self, i: int, val: int) -> None:
        self.update_recur(self.array, 0, len(self.array)-1, self.tree, 0, i, val)


    def update_recur(self, array, start, end, tree, inode, k, val):
        if len(array) > 0:
            if start == end:
                tree[inode] += val-array[k]
                array[k] = val
            else:
                mid = int((start+end)/2) 
                if  start<=k<=mid:
                    self.update_recur(array, start, mid, tree, inode*2+1, k, val)
                else:
                    self.update_recur(array, mid+1, end, tree, inode*2+2, k, val)
                tree[inode] = tree[inode*2+1] + tree[inode*2+2]
            return

    def sumRange(self, i: int, j: int) -> int:
        return self.subsum_recur(self.array, 0, len(self.array)-1, self.tree, 0, i, j)


    def subsum_recur(self, array, start, end, tree, inode, left, right):
        if len(array) == 0:
            return 
        else:
            if left<=start and right >= end:
                return tree[inode]
            if left>end or right<start:
                return 0
            else:
                mid = int((start+end)/2)
                sleft = self.subsum_recur(array, start, mid, tree, inode*2+1, left, right)
                sright = self.subsum_recur(array, mid+1, end, tree, inode*2+2, left, right)
                return sleft + sright

# Your NumArray object will be instantiated and called as such:
# obj = NumArray(nums)
# obj.update(i,val)
# param_2 = obj.sumRange(i,j)