LeetCode 筆記 - 307.Range Sum Query - Mutable

Posted on

題目在此 307. Range Sum Query - Mutable

給定一個整數陣列，實作一個可以更新數值並且可以計算區間和的資料結構。

解題思維

這題就是要你實作一個 Segment Tree (線段樹)，這其實算是一個比較進階的資料結構，主要是用來解決區間查詢和更新的問題。
LeetCode 上相關的問題可以參考 Segment Tree Problem List，可以看到超過三分之二都是困難題 😬

Segment Tree 的基本概念是將陣列分割成多個區間，並且將這些區間的和存儲在一棵樹狀結構中。這樣當我們需要查詢某個區間的和時，可以快速地從樹中找到對應的區間並計算出結果，而不需要每次都遍歷整個陣列。

Segment Tree 的實作通常包含以下幾個步驟：

  1. 建構樹：將陣列分割成多個區間，並且將這些區間的和存儲在樹的節點中。
  2. 更新數值：當陣列中的某個數值被更新時，需要更新樹中對應的區間和。
  3. 查詢區間和：當需要查詢某個區間的和時，可以快速地從樹中找到對應的區間並計算出結果。

而時間複雜度方面，建構樹的時間複雜度是 $O(n)$，更新數值和查詢區間和的時間複雜度都是 $O(\log n)$。
在面對大量的更新和查詢操作時，可以大幅提升效率。

程式碼

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
class NumArray:

    def __init__(self, nums: List[int]):
        self.nums = nums
        
        # 這兩種結構對速度沒有什麼差異。
        self.root = [0] * (len(nums) * 4)
        # self.root = collections.defaultdict(int)

        def build(node, index, left, right):
            if left == right:
                node[index] = nums[left]
                return

            mid = (left + right) // 2
            
            left_node = 2 * index + 1
            right_node = 2 * index + 2

            build(node, left_node, left, mid)
            build(node, right_node, mid + 1, right)

            node[index] = node[left_node] + node[right_node]

        build(self.root, 0, 0, len(nums) - 1)

    def update(self, index: int, val: int) -> None:
        self.nums[index] = val

        def update(node, cur_idx, left, right):
            if left == right:
                node[cur_idx] = val
                return

            mid = (left + right) // 2
            left_node = 2 * cur_idx + 1
            right_node = 2 * cur_idx + 2

            if index <= mid:
                update(node, left_node, left, mid)
            else:
                update(node, right_node, mid + 1, right)

            node[cur_idx] = node[left_node] + node[right_node]

        update(self.root, 0, 0, len(self.nums) - 1)

    def sumRange(self, left: int, right: int) -> int:

        def sumRange(node, cur_idx, cur_left, cur_right):
            if right < cur_left or cur_right < left:
                return 0

            if left <= cur_left and cur_right <= right:
                return node[cur_idx]

            mid = (cur_left + cur_right) // 2
            left_node = 2 * cur_idx + 1
            right_node = 2 * cur_idx + 2

            return sumRange(node, left_node, cur_left, mid) + \
                sumRange(node, right_node, mid + 1, cur_right)

        return sumRange(self.root, 0, 0, len(self.nums) - 1)

也許你也會想看看