2 Python solutions using Segment Tree.


  • 1
    M

    Array-based implementation (440ms):

    class NumArray(object):
        def __init__(self, nums):
            """
            initialize your data structure here.
            :type nums: List[int]
            """
            def makeTree(s_i, e_i, t_i):
                st = self.seg_tree
                if s_i == e_i:
                    st[t_i] = nums[s_i]
                    return st[t_i]
                mid = s_i + (e_i - s_i) // 2
                st[t_i] = makeTree(s_i, mid, 2 * t_i + 1) + makeTree(mid + 1, e_i, 2 * t_i + 2)
                return st[t_i]
    
            if nums:
                tree_length = 2 * (2 ** int(math.ceil(math.log(len(nums), 2)))) - 1
                self.nums = nums
                self.seg_tree = [0] * tree_length
                makeTree(0, len(nums) - 1, 0)
    
    
        def update(self, i, val):
            """
            :type i: int
            :type val: int
            :rtype: int
            """
            def update_until(s_i, e_i, i, t_i):
                if i > e_i or i < s_i:
                    return
                st[t_i] += diff
                if s_i == e_i:
                    return
                mid = s_i + (e_i - s_i) // 2
                update_until(s_i, mid, i, 2 * t_i + 1)
                update_until(mid + 1, e_i, i, 2 * t_i + 2)
    
            st = self.seg_tree
            diff = val - self.nums[i]
            self.nums[i] = val
            update_until(0, len(self.nums) - 1, i, 0)
    
    
        def sumRange(self, i, j):
            """
            sum of elements nums[i..j], inclusive.
            :type i: int
            :type j: int
            :rtype: int
            """
            def sumRange_until(s_i, e_i, t_i):
                if i <= s_i and j >= e_i:
                    return st[t_i]
                if s_i > j or e_i < i:
                    return 0
                mid = s_i + (e_i - s_i) // 2
                return sumRange_until(s_i, mid, 2 * t_i + 1) + sumRange_until(mid + 1, e_i, 2 * t_i + 2)
    
            st = self.seg_tree
            if i == j:
                return self.nums[i]
    
            return sumRange_until(0, len(self.nums) - 1, 0)
    

    Class-based implementation (540ms):

    class SegmentTreeNode(object):
        def __init__(self, val=0, start=0, end=0):
            self.val = val
            self.left = None
            self.right = None
            self.start = start
            self.end = end
    
    
    class NumArray(object):
        def __init__(self, nums):
            """
            initialize your data structure here.
            :type nums: List[int]
            """
            def make_tree(s_i, e_i):
                if s_i == e_i:
                    return SegmentTreeNode(nums[s_i], s_i, e_i)
                node = SegmentTreeNode()
                mid = s_i + (e_i - s_i) // 2
                node.left = make_tree(s_i, mid)
                node.right = make_tree(mid + 1, e_i)
                node.val = node.left.val + node.right.val
                node.start = s_i
                node.end = e_i
    
                return node
    
            if nums:
                self.nums = nums
                self.root = make_tree(0, len(nums) - 1)
    
        def update(self, i, val):
            """
            :type i: int
            :type val: int
            :rtype: int
            """
            def update_until(node):
                if not node or node.start > i or node.end < i:
                    return
                node.val += diff
                update_until(node.left)
                update_until(node.right)
            diff = val - self.nums[i]
            self.nums[i] = val
            update_until(self.root)
    
        def sumRange(self, i, j):
            """
            sum of elements nums[i..j], inclusive.
            :type i: int
            :type j: int
            :rtype: int
            """
            def sumRange_until(node):
                if not node or node.start > j or node.end < i:
                    return 0
                if i <= node.start and j >= node.end:
                    return node.val
                return sumRange_until(node.left) + sumRange_until(node.right)
    
            if i == j:
                return self.nums[i]
    
            return sumRange_until(self.root)

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.