307. Range Sum Query - Mutable

https://leetcode.com/problems/range-sum-query-mutable/

Python

Segment Tree

Space: O(N) # On the tree build time Time:

• Build: O(N)
• Update: O(logN)
• Query: O(logN)

from dataclasses import dataclass


@dataclass
class TreeNode:
    start: int
    end: int
    val: int = 0
    left: TreeNode = None
    right: TreeNode = None

class NumArray:
    def __init__(self, nums: List[int]):
        def build(start, end):
            if start == end:
                node = TreeNode(start, end, nums[start])
                return node
            middle = (start+end) >> 1
            left = build(start, middle)
            right = build(middle+1, end)
            node = TreeNode(
                start=start,
                end=end,
                val=left.val + right.val,
                left=left,
                right=right
            )
            return node
        self.tree = build(0, len(nums)-1)

    def update(self, index: int, val: int) -> None:
        def update_tree(node):
            if node.start == node.end == index:
                node.val = val
                return
            middle = (node.start+node.end) >> 1
            if index <= middle:
                update_tree(node.left)
            else:
                update_tree(node.right)
            node.val = node.left.val + node.right.val

        update_tree(self.tree)

    def sumRange(self, left: int, right: int) -> int:
        def query_tree(node, left, right):
            if node.start == left and node.end == right:
                return node.val

            middle = (node.start+node.end) >> 1

            if right <= middle:
                return query_tree(node.left, left, right)
            elif left > middle:
                return query_tree(node.right, left, right)
            else:
                return query_tree(node.left, left, middle) + query_tree(node.right, middle+1, right)

        return query_tree(self.tree, left, right)