|
| 1 | +# 307. Range Sum Query - Mutable |
| 2 | + |
| 3 | +## Segment Tree Solution |
| 4 | +- Run-time: create_tree() is O(N), sumRange() is O(logN), update() is O(logN) |
| 5 | +- Space: O(N) |
| 6 | +- N = Number of given nums |
| 7 | + |
| 8 | +A segment tree is similar to a binary tree, in fact, the only difference is that segment trees traversal via. the range of indexes. |
| 9 | +Each node of the tree, instead of storing a value, stores the values between the range of indexes. |
| 10 | +Segment trees are good if you have a lot of numbers and need to either find the max, min, sum, etc... of a given range. |
| 11 | + |
| 12 | +Naive methods would require O(N) time traversal to find the result of a given range. |
| 13 | +Other methods use a 2d array of O(N^2) space and O(N^2) run-time to pre-process every range but O(1) look up after. |
| 14 | +These methods don't work well when there are billions of numbers, therefore, segment trees are a great alternative. |
| 15 | + |
| 16 | +``` |
| 17 | +class NumArray: |
| 18 | + |
| 19 | + def __init__(self, nums: List[int]): |
| 20 | + self.tree = SegmentTree(nums) |
| 21 | + |
| 22 | + def update(self, i: int, val: int) -> None: |
| 23 | + self.tree.update(i, val) |
| 24 | + |
| 25 | + def sumRange(self, i: int, j: int) -> int: |
| 26 | + return self.tree.get_range(i, j) |
| 27 | + |
| 28 | +class SegmentTree(object): |
| 29 | + |
| 30 | + def __init__(self, nums): |
| 31 | + self.root = self._create_tree(nums) |
| 32 | + |
| 33 | + def _create_tree(self, nums): |
| 34 | + |
| 35 | + def tree_builder(left, right): |
| 36 | + if left > right: |
| 37 | + return None |
| 38 | + if left == right: |
| 39 | + return Node(nums[left], left, right) |
| 40 | + mid_idx = (left + right) // 2 |
| 41 | + left_node = tree_builder(left, mid_idx) |
| 42 | + right_node = tree_builder(mid_idx + 1, right) |
| 43 | + total_sum = left_node.sum if left_node is not None else 0 |
| 44 | + total_sum += right_node.sum if right_node is not None else 0 |
| 45 | + new_node = Node(total_sum, start=left, end=right, left=left_node, right=right_node) |
| 46 | + return new_node |
| 47 | + |
| 48 | + return tree_builder(0, len(nums) - 1) |
| 49 | + |
| 50 | + def update(self, idx, val): |
| 51 | + |
| 52 | + def update_helper(curr): |
| 53 | + if curr.start_idx == curr.end_idx == idx: # leaf |
| 54 | + curr.sum = val |
| 55 | + return |
| 56 | + mid_idx = (curr.start_idx + curr.end_idx) // 2 |
| 57 | + if idx <= mid_idx: # go left |
| 58 | + update_helper(curr.left) |
| 59 | + else: # go right |
| 60 | + update_helper(curr.right) |
| 61 | + curr.sum = curr.left.sum + curr.right.sum |
| 62 | + |
| 63 | + update_helper(self.root) |
| 64 | + |
| 65 | + def get_range(self, l, r): |
| 66 | + |
| 67 | + def sum_helper(curr, left, right): |
| 68 | + if left == curr.start_idx and right == curr.end_idx: # total overlap |
| 69 | + return curr.sum |
| 70 | + mid_idx = (curr.start_idx + curr.end_idx) // 2 |
| 71 | + if right <= mid_idx: # range is only on the left subtree? |
| 72 | + return sum_helper(curr.left, left, right) |
| 73 | + elif left >= mid_idx + 1: # range is only on the right subtree? |
| 74 | + return sum_helper(curr.right, left, right) |
| 75 | + # ranges are in both left and right subtrees |
| 76 | + return sum_helper(curr.left, left, mid_idx) + sum_helper(curr.right, mid_idx + 1, right) |
| 77 | + |
| 78 | + return sum_helper(self.root, l, r) |
| 79 | + |
| 80 | +class Node(object): |
| 81 | + |
| 82 | + def __init__(self, _sum, start, end, left=None, right=None): |
| 83 | + self.start_idx = start |
| 84 | + self.end_idx = end |
| 85 | + self.sum = _sum |
| 86 | + self.left = left |
| 87 | + self.right = right |
| 88 | +``` |
0 commit comments