區間演算法題用線段樹可以秒解？

背景

給一個兩個陣列，其中一個陣列是 A [1,2,3,4]，另外一個陣列是 B [5,6,7,8]。讓你求兩個陣列合並後的大陣列的：

• 最大值
• 最小值
• 總和

這題是不是很簡單？我們直接可以很輕鬆地在 $O(m+n)$ 的時間解決，其中 m 和 n 分別為陣列 A 和 B 的大小。

那如果我可以修改 A 和 B 的某些值，並且我要求很多次最大值，最小值和總和呢？

樸素的思路是原地修改陣列，然後 $O(m+n)$ 的時間重新計算。顯然這並沒有利用之前計算好的結果，效率是不高的。 那有沒有效率更高的做法？

有！線段樹就可以解決。

線段樹是什麼？

線段樹本質上就是一棵樹。更準確地說，它是一顆二叉樹，而且它是一顆平衡二叉樹。關於為什麼是平衡二叉樹，我們後面會講，這裡大家先有這樣一個認識。

雖然是一棵二叉樹，但是線段樹我們通常使用陣列來模擬樹結構，而不是傳統的定義 TreeNode 。

一方面是因為實現起來容易，另外一方面是因為線段樹其實是一顆完全二叉樹，因此使用陣列直接模擬會很高效。這裡的原因我已經在之前寫的堆專題中的二叉堆實現的時候中講過了，大家可以在我的公眾號《力扣加加》回覆堆獲取。

解決什麼問題

正如它的名字，線段樹和線段（區間）有關。線段樹的每一個樹節點其實都儲存了一個區間（段）的資訊。然後這些區間的資訊如果滿足一定的性質就可以用線段樹來提高效能。

那：

1. 究竟是什麼樣的性質？
2. 如何提高的效能呢？

究竟是什麼樣的性質？

比如前面我們提到的最大值，最小值以及求和就滿足這個一定性質。即我可以根據若干個（這裡是兩個）子集推匯出子集的並集的某一指標。

以上面的例子來說，我們可以將陣列 A 和 陣列 B 看成兩個集合。那麼：集合 A 的最大值和集合 B 的最大值已知，我們可以直接通過 max(Amax, Bmax) 求得集合 A 與集合 B 的並集的最大值。其中 Amax 和 Bmax 分別為集合 A 和集合 B 的最大值。最小值和總和也是一樣的，不再贅述。因此如果統計資訊滿足這種性質，我們就可以可以使用線段樹。但是要不要使用，還是要看用了線段樹後，是否能提高效能。

如何提高的效能呢？

關於提高效能，我先賣一個關子，等後面講完實現的時候，我們再聊。

實現

以文章開頭的求和為例。

我們可以將區間 A 和 區間 B 分別作為一個樹的左右節點，並將 A 的區間和與 B 的區間和分別儲存到左右子節點中。

接下來，將 A 的區間分為左右兩部分，同理 B 也分為左右兩部分。不斷執行此過程直到無法繼續分。

總結一下就是將區間不斷一分為二，並將區間資訊分別儲存到左右節點。如果是求和，那麼區間資訊就是區間的和。這個時候的線段樹大概是這樣的：

藍色字型表示的區間和。

注意，這棵樹的所有葉子節點一共有 n 個（n 為原陣列長度），並且每一個都對應到原陣列某一個值。

體現到程式碼上也很容易。 直接使用後續遍歷即可解決。這是因為，我們需要知道左右節點的統計資訊，才能計算出當前節點的統計資訊。

不熟悉後序遍歷的可以看下我之前的樹專題，公眾號力扣加加回復樹即可獲取

和二叉堆的表示方式一樣，我們可以用陣列表示樹，用 2 * i + 1 和 2 * 1 + 2 來表示左右節點的索引，其中 i 為當前節點對應在 tree 上的索引。

tree 是用來構建線段樹的陣列，和二叉堆類似。只不過 tree[i] 目前存的是區間資訊罷了。

上面我描述建樹的時候有明顯的遞迴性，因此我們可以遞迴的建樹。具體來說，可以定義一個 build(tree_index, l, r) 方法 來建樹。其中 l 和 r 就是對應區間的左右端點，這樣 l 和 r 就可以唯一確定一個區間。 tree_index 其實是用來標記當前的區間資訊應該被更新到 tree 陣列的哪個位置。

我們在 tree 上儲存區間資訊，那麼最終就可以用 tree[tree_index] = .... 來更新區間資訊啦。

核心程式碼：

def build(self, tree_index:int, l:int, r:int):
    '''
    遞迴建立線段樹
    tree_index : 線段樹節點在陣列中位置
    l, r : 該節點表示的區間的左，右邊界
    '''
    if l == r:
        self.tree[tree_index] = self.data[l]
        return
    mid = (l+r) // 2 # 區間中點，對應左孩子區間結束，右孩子區間開頭
    left, right = 2 * tree_index + 1, 2 * tree_index + 2 # tree_index 的左右子樹索引
    self.build(left, l, mid)
    self.build(right, mid+1, r)
    # 典型的後序遍歷
    # 區間和使用加法即可，如果不是區間和要改下面這行程式碼
    self.tree[tree_index] = self.tree[left] + self.tree[right]

上面程式碼的陣列 self.tree[i] 其實就是用來存類似上圖中藍色字型的區間和。每一個區間都在 tree 上存有它一個位置，存它的區間和。

複雜度分析

• 時間複雜度：由遞推關係式 T(n) = 2*T(n/2) + 1，因此時間複雜度為 $O(n)$

不知道怎麼得出的 $O(n)$? 可以看下我的《演算法通關之路》的第一章內容。 https://leetcode-solution.cn/...

• 空間複雜度：tree 的大小和 n 同階，因此空間複雜度為 $O(n)$

終於把樹建好了，但是知道現在一點都沒有高效起來。我們要做的是高效處理頻繁更新情況下的區間查詢。

那基於這種線段樹的方法，如果更新和查詢區間資訊如何做呢？

區間查詢

先回答簡單的問題區間查詢原理是什麼。

如果查詢一個區間的資訊。這裡也是使用後序遍歷就 ok 了。比如我要找一個區間 [l,r] 的區間和。

那麼如果當前左節點：

• 完整地落在 [l,r] 內。比如 [2,3] 完整地落在 [1,4] 內。 我們直接將 tree 中左節點對於的區間和取出來備用，不妨極為 lsum。
• 部分落在 [l,r] 內。比如 [1,3] 部分落在 [2,4]。這個時候我們繼續遞迴，直到完整地落在區間內（上面的那種情況），這個時候我們直接將 tree 中左節點對於的區間和取出來備用
• 將前面所有取出來備用的值加起來就是答案

右節點的處理也是一樣的，不再贅述。

複雜度分析

• 時間複雜度：查詢不需要在每個時刻都處理兩個葉子節點，實際上處理的次數大致和樹的高度一致。而樹是平衡的，因此複雜度為 $O(logn)$

或者由遞推關係式 T(n) = T(n/2) + 1，因此時間複雜度為 $O(logn)$

不知道怎麼得出的 $O(logn)$? 可以看下我的《演算法通關之路》的第一章內容。 https://leetcode-solution.cn/...

大家可以結合後面的程式碼理解這個複雜度。

區間修改

那麼如果我修改了 A[1] 為 1 呢？

如果不修改 tree，那麼顯然查詢的區間只要包含了 A[1] 就一定是錯的，比如查詢區間 [1,3] 的和 就會得到錯誤的答案。因此我們要在修改了 A[1] 的時候同時去修改 tree。

問題在於我們要修改哪些 tree 的值，修改為多少呢？

首先回答第一個問題，修改哪些 tree 的值呢？

我們知道，線段樹的葉子節點都是原陣列上的值，也是說，線段樹的 n 個葉子節點對應的就是原陣列。因此我們首先要找到我們修改的位置對應的那個葉子節點，將其值修改掉。

這就完了麼？

沒有完。實際上，我們修改的葉子節點的所有父節點以及祖父節點（如果有的話）都需要改。也就是說我們需要從這個葉子節點不斷冒泡到根節點，並修改沿途的區間資訊

這個過程和瀏覽器的事件模型是類似的

接下來回答最後一個問題，具體修改為多少？

對於求和，我們需要首先將葉子節點改為修改後的值，另外所有葉子節點到根節點路徑上的點的區間和都加上 delta，其中 delta 就是改變前後的差值。

求最大最小值如何更新？大家自己思考一下。

修改哪些節點，修改為多少的問題都解決了，那麼程式碼實現就容易了。

複雜度分析

• 時間複雜度：修改不需要在每個時刻都處理兩個葉子節點，實際上處理的次數大致和樹的高度一致。而樹是平衡的，因此複雜度為 $O(logn)$

或者由遞推關係式 T(n) = T(n/2) + 1，因此時間複雜度為 $O(logn)$

不知道怎麼得出的 $O(logn)$? 可以看下我的《演算法通關之路》的第一章內容。 https://leetcode-solution.cn/...

大家可以結合後面的程式碼理解這個複雜度。

程式碼

線段樹程式碼已經放在刷題外掛上了，公眾號《力扣加加》回覆外掛即可獲取。

class SegmentTree:
    def __init__(self, data:List[int]):
        '''
        data: 傳入的陣列
        '''
        self.data = data
        self.n = len(data)
        #  申請 4 倍 data 長度的空間來存線段樹節點
        self.tree = [None] * (4 * self.n) # 索引 i 的左孩子索引為 2i+1，右孩子為 2i+2
        if self.n:
            self.build(0, 0, self.n-1)
    # 本質就是一個自底向上的更新過程
    # 因此可以使用後序遍歷，即在函式返回的時候更新父節點。
    def update(self, tree_index, l, r, index):
        '''
        tree_index: 某個根節點索引
        l, r : 此根節點代表區間的左右邊界
        index : 更新的值的索引
        '''
        if l == r==index :
            self.tree[tree_index] = self.data[index]
            return
        mid = (l+r)//2
        left, right = 2 * tree_index + 1, 2 * tree_index + 2
        if index > mid:
            # 要更新的區間在右子樹
            self.update(right, mid+1, r, index)
        else:
            # 要更新的區間在左子樹 index<=mid
            self.update(left, l, mid, index)
        # 查詢區間一部分在左子樹一部分在右子樹
        # 區間和使用加法即可，如果不是區間和要改下面這行程式碼
        self.tree[tree_index] = self.tree[left] + self.tree[right]

    def updateSum(self,index:int,value:int):
        self.data[index] = value
        self.update(0, 0, self.n-1, index)
    def query(self, tree_index:int, l:int, r:int, ql:int, qr:int) -> int:
        '''
        遞迴查詢區間 [ql,..,qr] 的值
        tree_index : 某個根節點的索引
        l, r : 該節點表示的區間的左右邊界
        ql, qr: 待查詢區間的左右邊界
        '''
        if l == ql and r == qr:
            return self.tree[tree_index]

        # 區間中點，對應左孩子區間結束，右孩子區間開頭
        mid = (l+r) // 2
        left, right = tree_index * 2 + 1, tree_index * 2 + 2
        if qr <= mid:
            # 查詢區間全在左子樹
            return self.query(left, l, mid, ql, qr)
        elif ql > mid:
            # 查詢區間全在右子樹
            return self.query(right, mid+1, r, ql, qr)

        # 查詢區間一部分在左子樹一部分在右子樹
        # 區間和使用加法即可，如果不是區間和要改下面這行程式碼
        return self.query(left, l, mid, ql, mid) + self.query(right, mid+1, r, mid+1, qr)

    def querySum(self, ql:int, qr:int) -> int:
        '''
        返回區間 [ql,..,qr] 的和
        '''
        return self.query(0, 0, self.n-1, ql, qr)

    def build(self, tree_index:int, l:int, r:int):
        '''
        遞迴建立線段樹
        tree_index : 線段樹節點在陣列中位置
        l, r : 該節點表示的區間的左，右邊界
        '''
        if l == r:
            self.tree[tree_index] = self.data[l]
            return
        mid = (l+r) // 2 # 區間中點，對應左孩子區間結束，右孩子區間開頭
        left, right = 2 * tree_index + 1, 2 * tree_index + 2 # tree_index 的左右子樹索引
        self.build(left, l, mid)
        self.build(right, mid+1, r)
        # 區間和使用加法即可，如果不是區間和要改下面這行程式碼
        self.tree[tree_index] = self.tree[left] + self.tree[right]

相關專題

• 堆

大家可以看下我之前寫的堆的專題的二叉堆實現。然後對比學習，順便還學了堆，豈不美哉？

• 樹狀陣列

樹狀陣列和線段樹類似，難度比線段樹稍微高一點點。有機會給大家寫一篇樹狀陣列的文章。

• immutablejs

前端的小夥伴應該知道 immutable 吧？ 而 immutablejs 就是非常有名的實現 immutable 的工具庫。西法之前寫過一篇 immutable 原理解析文章，感興趣的可以看下 https://lucifer.ren/blog/2020...

回答前面的問題

為啥是平衡二叉樹？

前面的時間複雜度其實也是基於樹是平衡二叉樹這一前提。那麼線段樹一定是平衡二叉樹麼？是的。這是因為線段樹是完全二叉樹，而完全二叉樹是平衡的。

當然還有另外一個前提，那就是樹的總的節點數和原陣列長度同階，也就是 n 的量級。關於為啥是同階的，也容易計算，只需要根據遞迴公式即可得出。

為啥線段樹能提高效能？

只要你理解了我實現部分的時間複雜度，那麼就不難明白這個問題。因為修改和查詢的時間複雜度都是 $logn$，而不使用線段樹的暴力做法查詢的複雜度高達 $O(n)$。相應的代價就是建樹的 $O(n)$ 的空間，因此線段樹也是一種典型的空間換時間演算法。

最後點一下題。區間演算法題是否可以用線段樹秒解？這其實文章中已經回答過了，其取決於是否滿足兩點：

1. 是否可以根據若干個（這裡是兩個）子集推匯出子集的並集的某一指標。
2. 是否能提高效能（相比於樸素的暴力解法）。通常面臨頻繁修改的場景都可以考慮使用線段樹優化修改後的查詢時間消耗。