Optimizing Range Queries - Prefix Sum

Scenario#

通常在處理演算法問題，遇到題目給定陣列 nums，並針對區間和 [L, R] 範圍提問時，可能涵蓋以下幾種問題 :

區間查詢 (Range Query)
區間更新 (Range Update)
區間重疊 (Range Overlap)

三者中針對「區間查詢」且「靜態資料」的情況，正是適合 Prefix Sum 的解法。

靜態資料指的是資料不會變動。

Example#

以 Leetcode 3070. Count Submatrices with Top-Left Element and Sum Less Than k 為範例。

讓我們來了解一下題目限制：

包含左上角 (top-left) 區塊的 matrix
所有區塊加總需要 <= k
grid[r][c] >= 0

這題是屬於 2D matrix，因此如果要遍歷每個元件就需要 O(M * N)。

那如果直接用暴力解法的話，

第一層迴圈 : for 右下角的 row [0, m-1]
第二層迴圈 : for 右下角的 col [0, n-1]
第三、四層迴圈 : 為了算 (0,0) 到 (r,c) 的總和，要再跑兩個迴圈把裡面每一格加起來

從上面看到第一、二層迴圈，只是找到每個元素作為右下角 (r,c) 的可能時，才透過第三、四個迴圈計算總合。如此一來，時間複雜度就會是 O(M^2 * N^2)，非常地可觀。

Solution#

我們可以套用 Prefix Sum 與 Dynamic Programming 的概念，利用已經計算好的小矩陣，推導出大矩陣。

假設我們要計算到 (r,c) 的總和，可以從矩陣上拆成 :

左邊格子的總和 : (0,0) 到 (r,c-1)
上面格子的總和 : (0,0) 到 (r-1,c)
當前格子 : grid[r][c]

如果細心的讀者有發現，左邊與上面格子的加總會有一塊重複 (overlay) 的區塊，大小是 (0,0) 到 (r-1,c-1)。

因此最後得到的加總會是 Sum(r,c) = grid[r][c] + Sum(r-1,c) + Sum(r,c-1) - Sum(r-1,c-1)。

以下是最後的程式碼 :


class Solution:
    def countSubmatrices(self, grid: List[List[int]], k: int) -> int:
        m, n = len(grid), len(grid[0])
        ans = 0

        for r in range(m):
            for c in range(n):
                # 1. 取得「上面一格」的總和 (若 r=0 則為 0)
                up = grid[r-1][c] if r > 0 else 0

                # 2. 取得「左邊一格」的總和 (若 c=0 則為 0)
                left = grid[r][c-1] if c > 0 else 0

                # 3. 取得「重疊部分」的總和 (若 r=0 或 c=0 則為 0)
                overlap = grid[r-1][c-1] if (r > 0 and c > 0) else 0

                # 更新當前格子為「包含 (0,0) 的子矩陣總和」
                grid[r][c] = grid[r][c] + up + left - overlap

                # 4. 檢查這個子矩陣總和是否符合條件
                if grid[r][c] <= k:
                    ans += 1
                else:
                    # 由於已知 grid[r][c] >= 0，因此當 grid[r][c] 已經大於 k 之後
                    # 後面的數字只會越來越大，可以提早結束來進一步優化性能
                    break

        return ans

那最後的時間複雜度為 O(M*N)，空間複雜度為 O(1)，因為我們採用 in-place 修改。