前陣子把前綴和相關的題又系統(tǒng)過了一遍繞來繞去發(fā)現這套東西表面看著簡單真正上手寫的時候索引偏移、二維邊界、差分還原順序每一個都能讓人debug一晚上。前綴和Prefix Sum在Python里的實現往往只有幾行但它背后解決的是算法里非常典型的一類問題頻繁的區(qū)間查詢。如果你經常被子數組、子矩陣、區(qū)間求和這類題卡住或者在實際業(yè)務里需要對連續(xù)時間段的數據做累計統(tǒng)計這篇總結應該能幫你把前綴和徹底吃透。這篇文章面向的是已經會基礎Python語法、但還沒系統(tǒng)整理過前綴和技巧的讀者。我會從暴力解法為什么慢開始講然后給出一維、二維前綴和的完整實現再講差分數組這個前綴和的逆運算最后用幾道經典題走一遍從讀題到AC的完整過程。1. 從一道超時的題說起前綴和的動機與本質1.1 一個真實的暴力超時場景假設你現在拿到這樣一個需求有一個長度為n的數組比如[3, 1, 4, 1, 5, 9, 2, 6]然后有m次查詢每次問你某個區(qū)間[l, r]內的所有元素之和是多少。最直觀的寫法就是一個循環(huán)def range_sum(nums, queries): res [] for l, r in queries: total 0 for i in range(l, r 1): total nums[i] res.append(total) return res這個寫法沒毛病邏輯完全正確。但問題是假設n和m都是10的5次方量級每次查詢都要遍歷一遍區(qū)間最壞情況下總的時間復雜度是O(n*m)也就是10的10次方次操作。在Python里跑這個規(guī)?；镜扔诘人肋@就是典型的能跑但跑不動的代碼。我第一次在筆試里遇到這類題的時候就是老老實實寫了上面的暴力解法結果不出意外地超時了。那時候我才意識到面試官考察的根本不是你會不會寫循環(huán)求和而是你有沒有預處理的思維——把頻繁重復計算的東西提前算好。1.2 前綴和的數學本質前綴和的核心思想特別樸素我先準備一個數組pre里面存的是從數組開頭到當前位置的所有元素之和。比如對于上面的數組pre[0] 0一個方便計算的位置pre[1] 3pre[2] 3 1 4pre[3] 3 1 4 8...那么任意區(qū)間[l, r]的和可以直接用兩個前綴和相減得到pre[r1] - pre[l]。為什么是r1而不是r因為pre[i]的語義是前i個元素的和pre[l]意味著下標0到l-1這些元素的和pre[r1]意味著下標0到r這些元素的和兩者一減剩下的正好是下標l到r的元素。這個過程就是前綴和的核心數學本質區(qū)間和 兩個前綴和的差。預處理的時間是O(n)每次查詢的時間是O(1)總復雜度從O(n*m)降到了O(nm)。我自己的體會是前綴和本質上是在用空間換時間但換得非常劃算——它只需要一個長度等于n1的額外數組省下的卻是巨大量的重復求和運算。2. 一維前綴和公式、代碼與最容易踩的索引坑2.1 兩種常見的實現寫法一維前綴和的Python實現非常簡單常見的寫法有兩種。第一種是預先分配好長度的寫法def build_prefix(nums): n len(nums) pre [0] * (n 1) for i in range(1, n 1): pre[i] pre[i - 1] nums[i - 1] return pre第二種是直接用append動態(tài)構建def build_prefix(nums): pre [0] for x in nums: pre.append(pre[-1] x) return pre我推薦第二種寫法原因有兩點第一它在代碼上天然強調了pre[0] 0這個邊界第二它不容易出現下標錯位的問題因為你不需要在腦子里換算nums[i-1]還是nums[i]直接遍歷原始數組的每個元素就行。這里要強調一個關鍵的索引約定pre[i]表示的是前i個元素的和而不是下標i之前的和。這個約定貫穿所有前綴和的推導只要你在寫代碼時始終記得pre的長度是n1pre[i]對應nums[0..i-1]的和后面很多坑都能繞開。2.2 區(qū)間查詢的O(1)操作有了pre數組之后查詢區(qū)間[l, r]的和就變成了一行代碼def query(pre, l, r): return pre[r 1] - pre[l]比如查詢上面數組的[2, 5]區(qū)間對應的元素是4, 1, 5, 9和是19。用pre算pre[6] - pre[2]pre[6]是前6個元素31415923pre[2]是前2個元素314相減正好是19。這個操作沒有任何循環(huán)也不涉及任何乘法就是一個減法。在實際刷題的時候這種查詢之間互相獨立、查詢次數又多的場景前綴和幾乎是唯一的最優(yōu)解。我還見過有人把pre數組本身當作結果輸出然后查詢的時候寫pre[r] - pre[l-1]。這也能用但屬于另一種約定——pre[i]表示前i1個元素的和。不建議混用否則代碼里一會兒減一一會兒不減一非常容易出現思維混亂。2.3 索引偏移的經典錯誤與規(guī)避我當初學前綴和的時候寫過一段讓我極其痛苦的代碼pre [0] * n for i in range(1, n): pre[i] pre[i - 1] nums[i]這段代碼的問題是pre[0]一直等于0pre[1]等于nums[1]但nums[1]其實是數組的第二個元素。這樣一來所有從那個位置取值的區(qū)間和都會錯位一位。更隱蔽的是如果數組元素恰好有正有負某些查詢看起來結果又碰巧是對的導致我復盤的時候完全找不到問題在哪。規(guī)避這種索引坑的方法其實很簡單嚴格遵循偏移1位的約定讓pre[i]對應前i個元素的和。具體來說就是構建時pre[i] pre[i-1] nums[i-1]查詢時區(qū)間[l, r]的和 pre[r1] - pre[l]只要記住pre比nums長一位開頭多一個0你就再也不會被索引困擾了。另外如果你用的是append寫法幾乎天然就符合這個約定這也是我強烈推薦它的原因。3. 二維前綴和矩陣問題里的容斥原理3.1 從一維到二維的推廣一維前綴和解決的是數組區(qū)間問題二維前綴和解決的是矩陣子矩陣問題。思路是一樣的預處理一個和矩陣同形的二維數組每個位置存的是從左上角到當前位置這個矩形區(qū)域內的所有元素之和。假設原始矩陣是matrix行數和列數分別是m和n那么二維前綴和S可以這樣構建這里同樣采用偏移1位的寫法def build_2d_prefix(matrix): m len(matrix) n len(matrix[0]) S [[0] * (n 1) for _ in range(m 1)] for i in range(1, m 1): row_sum 0 for j in range(1, n 1): row_sum matrix[i - 1][j - 1] S[i][j] S[i - 1][j] row_sum return S我這里用了每行滾動累加的方式實際計算每個位置時S[i][j]等于上面一行同列的前綴和S[i-1][j]再加上本行從開頭到當前列的所有元素之和row_sum。這樣的寫法避免了重復計算整行的和效率會高一點。更常見的遞推公式是容斥形式S[i][j] S[i-1][j] S[i][j-1] - S[i-1][j-1] matrix[i-1][j-1]翻譯成直覺語言當前位置的值 上方矩形的和 左方矩形的和 - 左上角被重復算了一次的矩形的和 原矩陣當前位置的值。3.2 子矩陣和的查詢公式有了二維前綴和矩陣S要查詢任意一個左上角為(r1, c1)、右下角為(r2, c2)的子矩陣的元素和用的是下面這個公式def query_2d(S, r1, c1, r2, c2): return S[r2 1][c2 1] - S[r1][c2 1] - S[r2 1][c1] S[r1][c1]原因還是容斥原理S[r21][c21]是整個從左上角到(r2, c2)的矩形和先減去上半部分S[r1][c21]再減去左半部分S[r21][c1]這時候左上角被減了兩次所以再加回來一次S[r1][c1]。我建議在理解這個公式時畫一個4x4的小矩陣把每個S[i][j]都標注成從左上角開始的一塊區(qū)域然后實際算一次。這個方法我每次都推薦給入門的同學因為它比死記公式牢靠得多——你一旦畫過一遍容斥的加減邏輯就刻在腦子里了。3.3 二維場景的邊界與空矩陣處理二維前綴和最容易被忽略的坑是空矩陣的邊界判斷。如果matrix本身是空的或者matrix[0]是空的len(matrix[0])會直接報錯IndexError。所以在構建二維前綴和的入口處一定要記得加判斷if not matrix or not matrix[0]: return []另一個坑是查詢坐標的合法性。雖然通常題目保證查詢坐標合法但如果你寫的是供內部使用的工具函數最好自己加一層斷言或者防御性處理。我之前在線筆試的時候有一道題就是因為查詢坐標可能越界我沒有處理導致好幾個case直接RTE。還有一點值得注意二維前綴和矩陣是(m1) * (n1)的形狀比原始矩陣多一行一列。如果你習慣用numpy也可以直接用np.cumsum(matrix, axis0)再對行做一次cumsum得到相同效果但純Python實現時還是老老實實寫循環(huán)最穩(wěn)妥因為numpy的計算結果類型是ndarray在普通算法題里反倒要額外轉回list不算劃算。4. 差分數組前綴和的逆運算4.1 差分的核心思想如果說前綴和解決的核心問題是頻繁查詢區(qū)間和那差分數組解決的核心問題就是頻繁對區(qū)間進行整體加減操作最后一次查詢所有結果。什么叫區(qū)間整體加減舉個例子有一個數組現在給你一系列操作每個操作說從下標l到下標r的所有元素都加上一個值v。操作次數很多每個操作都去遍歷區(qū)間那時間復雜度自然又是O(n*m)。差分數組就是為了解決這個問題出現的。差分數組的定義是這樣的d [0] * n d[0] nums[0] for i in range(1, n): d[i] nums[i] - nums[i - 1]也就是說差分數組d[i]存的是原數組相鄰元素的差。這里最關鍵的性質是對差分數組求前綴和可以得到原來的數組。這正是差分是前綴和的逆運算這句話的含義。4.2 區(qū)間加法的O(1)實現當我們需要對區(qū)間[l, r]內的所有元素都加上一個值v時不需要去改原數組只需要在差分數組上操作兩個位置d[l] v if r 1 n: d[r 1] - v為什么這樣有效因為差分數組的語義是相鄰元素的差值。d[l] v會讓從l開始的所有元素在原數組上都多出v而d[r1] - v會把超出r這個范圍的增量抵消掉。這樣一來原數組從l到r的值都會加v而r1及之后的值不變。等所有操作都執(zhí)行完之后只需要對差分數組求一次前綴和就能還原出最終的原數組for i in range(1, n): d[i] d[i - 1]我看過很多初學者在這里犯迷糊其實可以把這個過程類比成在高鐵上給某一段車廂發(fā)紀念品你只需要在起點站上車發(fā)然后在終點站下車前把發(fā)放狀態(tài)取消掉不需要每站挨個發(fā)。差分就是這種狀態(tài)的記錄。4.3 什么時候用差分而不是線段樹我經常被問到一個問題區(qū)間加法和區(qū)間查詢都有很多次到底該用差分還是線段樹這里有個很實用的判斷標準如果操作是先批量修改、后統(tǒng)一查詢差分數組就是最優(yōu)解代碼簡單時間O(n)。如果操作是修改和查詢穿插進行比如改一次查一次再改再查那就需要線段樹或者樹狀數組了。如果每次查詢還需要實時返回某個區(qū)間的和并且在兩次查詢之間還有更新這種動態(tài)場景前綴和和差分都搞不定。換句話說差分數組和前綴和適合的是靜態(tài)場景數據在預處理階段就都定好了之后只是不停地查。如果題目里出現了在線這種詞往往就意味著要上數據結構的進階方案了。5. 前綴和的進階變體哈希優(yōu)化、異或前綴與滑動窗口的配合5.1 前綴和 哈希表解決子數組計數問題基本的前綴和能解決求某個區(qū)間的和但有一類更刁鉆的問題問的是有多少個子數組的和等于K。比如LeetCode 560題就是經典代表。如果還是用普通前綴和你可以先算出pre數組然后枚舉所有可能的左右端點檢查pre[r1] - pre[l] k這樣時間復雜度是O(n^2)。在數據量大的時候依然會掛。優(yōu)化思路是這樣的我們遍歷原數組計算當前位置的前綴和cur。對于每個cur我們需要知道的是在它之前有多少個前綴和的值等于cur - k。因為這些前綴和對應的位置正好是能讓cur - pre[i] k成立的位置。于是可以用一個哈希表來記錄每個前綴和值出現的次數def subarray_sum(nums, k): count {0: 1} cur 0 ans 0 for x in nums: cur x ans count.get(cur - k, 0) count[cur] count.get(cur, 0) 1 return ans這里有個很關鍵的細節(jié)初始化哈希表時要放入{0: 1}這代表一個空前綴和它的意義是如果當前的前綴和本身恰好等于k那么它和空前綴之間就形成了一個合法子數組。我自己第一次寫的時候漏了這個初始化結果所有case都少算了一部分答案。5.2 異或前綴和前綴和不止適用于加法對于異或運算同樣適用而且性質更好。我們把pre[i]定義成前i個元素的異或結果那么區(qū)間[l, r]的異或值就是pre[r1] ^ pre[l]。這里不需要容斥因為異或運算有一個自反性質x ^ x 0。所以兩個相同的前綴異或值相遇異或結果就是0中間的部分自然就露出來了。這個技巧在解決找出數組中所有異或和為0的子數組個數這類問題上非常有用。舉個例子LeetCode 525題連續(xù)數組這道題要找的是最長的連續(xù)子數組使得子數組中0和1的數量相同。一個很巧妙的做法是把0看成-1把1看成1用前綴和來記錄遍歷到當前位置時的累計和。當兩個位置的前綴和值相等時說明這一段的0和1數量相同因為它們的差值抵消了。5.3 前綴和與滑動窗口的邊界很多人會混淆滑動窗口和前前綴和的適用場景這里我?guī)湍惴智宄绻麛到M元素全為正數要求最短/最長滿足某個條件的子數組滑動窗口是最優(yōu)解因為窗口擴大或縮小的單調性讓雙指針能在線性時間內移動。如果數組元素有正有負滑動窗口的單調性失效因為窗口變長不一定讓和變大這時候前綴和往往配合哈希表或排序才能勝任。我曾經在一道題里先用了滑動窗口結果因為數組里有負數窗口的收縮邏輯陷入了死循環(huán)。后來改成前綴和加二分才解決。所以我個人建議遇到子數組和等于/大于/小于某個值這類問題第一時間先判斷數組里是否有負數。有負數基本就要往前綴和的方向思考了。6. Python實現中的性能與代碼風格建議6.1 用itertools.accumulate精簡代碼Python標準庫里的itertools.accumulate可以直接生成前綴和序列代碼會非常簡潔from itertools import accumulate nums [3, 1, 4, 1, 5, 9, 2, 6] pre [0] list(accumulate(nums))這個寫法的時間復雜度同樣是O(n)但代碼量幾乎降到了最低。我用這個方式在六十幾行的代碼里實現了一個小的數據統(tǒng)計腳本用于生成每日訂單的累計量曲線效果非常理想。不過要注意的是accumulate返回的是一個迭代器如果你需要隨機訪問某個位置的前綴和必須通過list()轉成列表否則沒法按下標取值。6.2 大數據量下的Python取舍前綴和本身的時間復雜度是O(n)但Python在大規(guī)模數據下會暴露出語言性能的天花板。比如處理10的7次方級別的數據時append循環(huán)的速度會比C的同等方式慢很多。我在實際處理千萬級流量日志的累計分布時最后是切到了numpy.cumsum速度一下子提升了十幾倍import numpy as np arr np.array(nums) pre_arr np.cumsum(arr)所以我的建議是算法刷題場景下用純Python的append寫法就夠這樣能保持代碼的可讀性也避免引入numpy后在線評測系統(tǒng)上因庫缺失或版本問題報錯。但如果是本地處理真實業(yè)務數據、數據量大且有numpy環(huán)境完全可以放開用np.cumsum。另外Python的整數可以無限大計算前綴和時不會像C那樣有int溢出的問題這算是Python一個隱形的福利在累加大量正整數時特別省心。6.3 刷題習慣與邊界測試我見過很多人在刷前綴和題目時代碼測了幾個樣例能過就急著交結果一提交就WA。這往往不是因為思路錯了而是邊界沒測到。我總結了一份前綴和必測的邊界清單數組長度為0或1的情況查詢的l等于0的情況這最考驗前綴和pre[0] 0的約定是否準確查詢區(qū)間覆蓋整個數組的情況數組中有大量0或者全是負數的情況差分數組操作中區(qū)間右端點恰好等于n-1的情況此時r1正好越界每條邊界都可以用幾行小樣例驗證別嫌麻煩。我自己的做法是寫一個簡單的暴力解法和前綴和解法對照著跑隨機數據兩邊結果不一致時再debug效率高很多。這個方法也推薦給你等于用一個笨辦法來給聰明辦法兜底。7. 實戰(zhàn)驗證四道經典題從讀題到AC的完整走查7.1 LeetCode 303區(qū)域和檢索 - 數組不可變這道題幾乎是純前綴和的入門題。題目的要求是構建一個類支持sumRange(left, right)方法返回從left到right的元素和。我的實現思路是直接在初始化時構建好前綴和數組然后sumRange里做一次減法class NumArray: def __init__(self, nums): self.pre [0] for x in nums: self.pre.append(self.pre[-1] x) def sumRange(self, left: int, right: int) - int: return self.pre[right 1] - self.pre[left]這道題我推薦所有人親手寫一遍因為它是所有前綴和題型的骨架。你把它寫熟了后面的二維和差分題都會順暢很多。初始化O(n)查詢O(1)這個復雜度就是前綴和的標準表現。7.2 LeetCode 304二維區(qū)域和檢索 - 矩陣不可變這是303的二維升級版。構建二維前綴和矩陣后查詢時用容斥公式邏輯和前面第3章講的一致class NumMatrix: def __init__(self, matrix): if not matrix or not matrix[0]: self.S [] return m, n len(matrix), len(matrix[0]) S [[0] * (n 1) for _ in range(m 1)] for i in range(1, m 1): for j in range(1, n 1): S[i][j] (S[i-1][j] S[i][j-1] - S[i-1][j-1] matrix[i-1][j-1]) self.S S def sumRegion(self, row1, col1, row2, col2): S self.S return (S[row21][col21] - S[row1][col21] - S[row21][col1] S[row1][col1])我在這道題的實現上犯過一個迷糊構建時把S[i][j]的遞推關系寫成了S[i-1][j-1] matrix[i-1][j-1]漏掉了左方和上方那兩個大矩形結果查詢小于3x3的矩陣時是對的一旦矩陣變大就全亂了。所以構建二維前綴和時一定要想清楚你算的究竟是整個左上矩形還是某個小矩形別貪圖少寫一個變量。7.3 LeetCode 1109航班預訂統(tǒng)計差分經典題這道題描述是這樣的有n個航班用1到n編號現在有bookings[i] [first_i, last_i, seats_i]表示從first_i到last_i的航班每個都預定了seats_i個座位最后返回每個航班的總預定數。這題就是典型的多次區(qū)間更新最后一次性查詢直接用差分數組class Solution: def corpFlightBookings(self, bookings, n): diff [0] * (n 1) for l, r, seats in bookings: diff[l - 1] seats diff[r] - seats ans [] cur 0 for i in range(n): cur diff[i] ans.append(cur) return ans很多人在這道題里卡住的一個點是題目里的航班編號是1-indexed而數組下標是0-indexed所以first_i要減1再對應到差分數組下標。同時差分數組我開到了n1的長度這樣即使r恰好等于n也不會越界。這道題如果用暴力每次預訂都要遍歷一遍區(qū)間復雜度是O(n * bookings)大概率超時。而差分數組的解法只需要O(len(bookings) n)幾乎是一邊遍歷一邊就出結果了。我強烈建議你用這道題來檢驗自己對差分理解的深度因為面試里它出現的頻率很高而且換個馬甲就是考同一套東西。7.4 LeetCode 560和為 K 的子數組前綴和哈希這道題我在第5節(jié)已經給出核心代碼這里再補一個完整的走查過程。題目問的是連續(xù)子數組中和為k的個數注意這里子數組必須是連續(xù)的而且元素可能有負數所以滑動窗口直接出局。我的思路是這樣的從左到右遍歷數組維護當前前綴和cur同時維護一個哈希表記錄每個前綴和值出現的次數。對于當前位置cur - k這個值如果在之前的某個前綴和位置出現過那么就說明那一段到當前這一段的和恰好是k。累加這些次數就是一個合法的答案數。class Solution: def subarraySum(self, nums, k): prefix_count {0: 1} cur 0 ans 0 for x in nums: cur x ans prefix_count.get(cur - k, 0) prefix_count[cur] prefix_count.get(cur, 0) 1 return ans這里最關鍵的一步是先查答案再更新哈希表。如果先更新哈希表再查就會把當前這個位置自己也算進去導致重復計數。我實際調試這道題時就是因為把這兩行順序寫反了答案始終比期望值大了一圈。后來我把中間過程打出來才意識到當前前綴和不能和它自己匹配。這道題吃透之后你再看一些類似的子數組和統(tǒng)計題目會發(fā)現套路高度一致哈希表存歷史前綴和遍歷時做差查表再更新當前前綴和。這就是一法通、萬法通的效果。我自己實際用Python跑過這四道題的完整流程從讀題到AC基本都在五分鐘以內核心的思考時間幾乎全部花在確認這道題是前綴和/差分場景還是需要更復雜的數據結構上。等你把前綴和的幾個變體都練熟了這種判斷就會變成一種直覺——看到區(qū)間查詢想前綴和看到批量區(qū)間更新想差分看到子數組計數想哈希表輔助一秒鐘就能定下方向。