Interview problem: Segment Tree

I heard the name of segment tree several times before and never had an opportunity to learn what it is. Today I finally understand the concepts of segment tree by watching these seven short youtube videos 9.1-9.7. These lectures are highly recommended to those who are new to the topics of segment tree like me.

Here is my python solution to the Leetcode 307 using the segment tree:

class TreeNode:
    def __init__(self, start, end, total, left=None, right=None):
        self.start = start
        self.end = end
        self.total = total
        self.left = left
        self.right = right


class NumArray:
    def __init__(self, nums: List[int]):
        def _create(start, end):
            if start == end:
                return TreeNode(start, end, nums[start])
            mid = start + (end - start) // 2
            left = _create(start, mid)
            right = _create(mid + 1, end)
            return TreeNode(start, end, left.total + right.total, left, right)
        self.root = _create(0, len(nums) - 1)

    def update(self, index: int, val: int) -> None:
        def _update(root):
            if root.start == root.end:
                root.total = val
                return
            mid = root.start + (root.end - root.start) // 2
            if index <= mid:
                _update(root.left)
            else:
                _update(root.right)
            root.total = root.left.total + root.right.total
        return _update(self.root)

    def sumRange(self, left: int, right: int) -> int:
        def _query(root, start, end):
            if start == root.start and end == root.end:
                return root.total
            mid = root.start + (root.end - root.start) // 2
            if end <= mid:
                return _query(root.left, start, end)
            if start >= mid + 1:
                return _query(root.right, start, end)
            return _query(root.left, start, mid) + _query(root.right, mid + 1, end)
        return _query(self.root, left, right)

Comments

Popular posts from this blog

529 Plan

How to offset W2 tax

Health Saving Account (HSA)