前阵子把前缀和相关的题又系统过了一遍,绕来绕去发现这套东西表面看着简单,真正上手写的时候,索引偏移、二维边界、差分还原顺序,每一个都能让人debug一晚上。前缀和(Prefix Sum)在Python里的实现往往只有几行,但它背后解决的是算法里非常典型的一类问题:频繁的区间查询。如果你经常被子数组、子矩阵、区间求和这类题卡住,或者在实际业务里需要对连续时间段的数据做累计统计,这篇总结应该能帮你把前缀和彻底吃透。
这篇文章面向的是已经会基础Python语法、但还没系统整理过前缀和技巧的读者。我会从暴力解法为什么慢开始讲,然后给出一维、二维前缀和的完整实现,再讲差分数组这个"前缀和的逆运算",最后用几道经典题走一遍从读题到AC的完整过程。
1. 从一道超时的题说起:前缀和的动机与本质
1.1 一个真实的暴力超时场景
假设你现在拿到这样一个需求:有一个长度为n的数组,比如[3, 1, 4, 1, 5, 9, 2, 6],然后有m次查询,每次问你某个区间[l, r]内的所有元素之和是多少。
最直观的写法就是一个循环:
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次方量级,每次查询都要遍历一遍区间,最坏情况下总的时间复杂度是O(n*m),也就是10的10次方次操作。在Python里跑这个规模基本等于等死,这就是典型的"能跑但跑不动"的代码。
我第一次在笔试里遇到这类题的时候,就是老老实实写了上面的暴力解法,结果不出意外地超时了。那时候我才意识到,面试官考察的根本不是你会不会写循环求和,而是你有没有"预处理"的思维——把频繁重复计算的东西提前算好。
1.2 前缀和的数学本质
前缀和的核心思想特别朴素:我先准备一个数组pre,里面存的是"从数组开头到当前位置的所有元素之和"。比如对于上面的数组:
pre[0] = 0(一个方便计算的位置)pre[1] = 3pre[2] = 3 + 1 = 4pre[3] = 3 + 1 + 4 = 8- ...
那么任意区间[l, r]的和,可以直接用两个前缀和相减得到:pre[r+1] - pre[l]。
为什么是r+1而不是r?因为pre[i]的语义是"前i个元素的和",pre[l]意味着下标0到l-1这些元素的和,pre[r+1]意味着下标0到r这些元素的和,两者一减,剩下的正好是下标l到r的元素。
这个过程就是前缀和的核心数学本质:区间和 = 两个前缀和的差。预处理的时间是O(n),每次查询的时间是O(1),总复杂度从O(n*m)降到了O(n+m)。
我自己的体会是:前缀和本质上是在用空间换时间,但换得非常划算——它只需要一个长度等于n+1的额外数组,省下的却是巨大量的重复求和运算。
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动态构建:
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的长度是n+1,pre[i]对应nums[0..i-1]的和",后面很多坑都能绕开。
2.2 区间查询的O(1)操作
有了pre数组之后,查询区间[l, r]的和就变成了一行代码:
def query(pre, l, r): return pre[r + 1] - pre[l]比如查询上面数组的[2, 5]区间,对应的元素是4, 1, 5, 9,和是19。用pre算:pre[6] - pre[2],pre[6]是前6个元素3+1+4+1+5+9=23,pre[2]是前2个元素3+1=4,相减正好是19。
这个操作没有任何循环,也不涉及任何乘法,就是一个减法。在实际刷题的时候,这种"查询之间互相独立、查询次数又多"的场景,前缀和几乎是唯一的最优解。
我还见过有人把pre数组本身当作结果输出,然后查询的时候写pre[r] - pre[l-1]。这也能用,但属于另一种约定——pre[i]表示前i+1个元素的和。不建议混用,否则代码里一会儿减一一会儿不减一,非常容易出现思维混乱。
2.3 索引偏移的经典错误与规避
我当初学前缀和的时候,写过一段让我极其痛苦的代码:
pre = [0] * n for i in range(1, n): pre[i] = pre[i - 1] + nums[i]这段代码的问题是:pre[0]一直等于0,pre[1]等于nums[1],但nums[1]其实是数组的第二个元素。这样一来,所有从那个位置取值的区间和都会错位一位。更隐蔽的是,如果数组元素恰好有正有负,某些查询看起来结果又"碰巧"是对的,导致我复盘的时候完全找不到问题在哪。
规避这种索引坑的方法其实很简单:严格遵循"偏移1位"的约定,让pre[i]对应前i个元素的和。具体来说就是:
- 构建时:
pre[i] = pre[i-1] + nums[i-1] - 查询时:区间
[l, r]的和 =pre[r+1] - pre[l]
只要记住"pre比nums长一位,开头多一个0",你就再也不会被索引困扰了。另外,如果你用的是append写法,几乎天然就符合这个约定,这也是我强烈推荐它的原因。
3. 二维前缀和:矩阵问题里的容斥原理
3.1 从一维到二维的推广
一维前缀和解决的是数组区间问题,二维前缀和解决的是矩阵子矩阵问题。思路是一样的:预处理一个和矩阵同形的二维数组,每个位置存的是"从左上角到当前位置这个矩形区域内的所有元素之和"。
假设原始矩阵是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[r2+1][c2+1]是整个从左上角到(r2, c2)的矩形和,先减去上半部分S[r1][c2+1],再减去左半部分S[r2+1][c1],这时候左上角被减了两次,所以再加回来一次S[r1][c1]。
我建议在理解这个公式时画一个4x4的小矩阵,把每个S[i][j]都标注成从左上角开始的一块区域,然后实际算一次。这个方法我每次都推荐给入门的同学,因为它比死记公式牢靠得多——你一旦画过一遍,容斥的加减逻辑就刻在脑子里了。
3.3 二维场景的边界与空矩阵处理
二维前缀和最容易被忽略的坑是空矩阵的边界判断。如果matrix本身是空的,或者matrix[0]是空的,len(matrix[0])会直接报错IndexError。所以在构建二维前缀和的入口处,一定要记得加判断:
if not matrix or not matrix[0]: return []另一个坑是查询坐标的合法性。虽然通常题目保证查询坐标合法,但如果你写的是供内部使用的工具函数,最好自己加一层断言或者防御性处理。我之前在线笔试的时候,有一道题就是因为查询坐标可能越界我没有处理,导致好几个case直接RTE。
还有一点值得注意:二维前缀和矩阵是(m+1) * (n+1)的形状,比原始矩阵多一行一列。如果你习惯用numpy,也可以直接用np.cumsum(matrix, axis=0)再对行做一次cumsum得到相同效果,但纯Python实现时还是老老实实写循环最稳妥,因为numpy的计算结果类型是ndarray,在普通算法题里反倒要额外转回list,不算划算。
4. 差分数组:前缀和的逆运算
4.1 差分的核心思想
如果说前缀和解决的核心问题是"频繁查询区间和",那差分数组解决的核心问题就是"频繁对区间进行整体加减操作,最后一次查询所有结果"。
什么叫区间整体加减?举个例子:有一个数组,现在给你一系列操作,每个操作说"从下标l到下标r的所有元素都加上一个值v"。操作次数很多,每个操作都去遍历区间,那时间复杂度自然又是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 区间加法的O(1)实现
当我们需要对区间[l, r]内的所有元素都加上一个值v时,不需要去改原数组,只需要在差分数组上操作两个位置:
d[l] += v if r + 1 < n: d[r + 1] -= v为什么这样有效?因为差分数组的语义是"相邻元素的差值"。d[l] += v会让从l开始的所有元素在原数组上都多出v,而d[r+1] -= v会把超出r这个范围的增量抵消掉。这样一来,原数组从l到r的值都会加v,而r+1及之后的值不变。
等所有操作都执行完之后,只需要对差分数组求一次前缀和,就能还原出最终的原数组:
for i in range(1, n): d[i] += d[i - 1]我看过很多初学者在这里犯迷糊,其实可以把这个过程类比成"在高铁上给某一段车厢发纪念品":你只需要在起点站上车发,然后在终点站下车前把"发放状态"取消掉,不需要每站挨个发。差分就是这种状态的记录。
4.3 什么时候用差分而不是线段树
我经常被问到一个问题:区间加法和区间查询都有很多次,到底该用差分还是线段树?这里有个很实用的判断标准:
- 如果操作是先批量修改、后统一查询,差分数组就是最优解,代码简单,时间O(n)。
- 如果操作是修改和查询穿插进行,比如改一次查一次再改再查,那就需要线段树或者树状数组了。
- 如果每次查询还需要实时返回某个区间的和,并且在两次查询之间还有更新,这种动态场景前缀和和差分都搞不定。
换句话说,差分数组和前缀和适合的是"静态"场景:数据在预处理阶段就都定好了,之后只是不停地查。如果题目里出现了"在线"这种词,往往就意味着要上数据结构的进阶方案了。
5. 前缀和的进阶变体:哈希优化、异或前缀与滑动窗口的配合
5.1 前缀和 + 哈希表:解决子数组计数问题
基本的前缀和能解决"求某个区间的和",但有一类更刁钻的问题问的是"有多少个子数组的和等于K"。比如LeetCode 560题就是经典代表。
如果还是用普通前缀和,你可以先算出pre数组,然后枚举所有可能的左右端点,检查pre[r+1] - pre[l] == k,这样时间复杂度是O(n^2)。在数据量大的时候依然会挂。
优化思路是这样的:我们遍历原数组,计算当前位置的前缀和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这里有个很关键的细节:初始化哈希表时要放入{0: 1},这代表"一个空前缀和",它的意义是如果当前的前缀和本身恰好等于k,那么它和空前缀之间就形成了一个合法子数组。我自己第一次写的时候漏了这个初始化,结果所有case都少算了一部分答案。
5.2 异或前缀和
前缀和不止适用于加法,对于异或运算同样适用,而且性质更好。我们把pre[i]定义成前i个元素的异或结果,那么区间[l, r]的异或值就是pre[r+1] ^ pre[l]。
这里不需要容斥,因为异或运算有一个自反性质:x ^ x = 0。所以两个相同的前缀异或值相遇,异或结果就是0,中间的部分自然就露出来了。这个技巧在解决"找出数组中所有异或和为0的子数组个数"这类问题上非常有用。
举个例子,LeetCode 525题"连续数组",这道题要找的是最长的连续子数组,使得子数组中0和1的数量相同。一个很巧妙的做法是:把0看成-1,把1看成+1,用前缀和来记录遍历到当前位置时的累计和。当两个位置的前缀和值相等时,说明这一段的0和1数量相同,因为它们的差值抵消了。
5.3 前缀和与滑动窗口的边界
很多人会混淆滑动窗口和前前缀和的适用场景,这里我帮你分清楚:
- 如果数组元素全为正数,要求最短/最长满足某个条件的子数组,滑动窗口是最优解,因为窗口扩大或缩小的单调性让双指针能在线性时间内移动。
- 如果数组元素有正有负,滑动窗口的单调性失效,因为窗口变长不一定让和变大,这时候前缀和(往往配合哈希表或排序)才能胜任。
我曾经在一道题里先用了滑动窗口,结果因为数组里有负数,窗口的收缩逻辑陷入了死循环。后来改成前缀和加二分才解决。所以我个人建议:遇到"子数组和等于/大于/小于某个值"这类问题,第一时间先判断数组里是否有负数。有负数,基本就要往前缀和的方向思考了。
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),但代码量几乎降到了最低。我用这个方式在六十几行的代码里实现了一个小的数据统计脚本,用于生成每日订单的累计量曲线,效果非常理想。
不过要注意的是,accumulate返回的是一个迭代器,如果你需要随机访问某个位置的前缀和,必须通过list()转成列表,否则没法按下标取值。
6.2 大数据量下的Python取舍
前缀和本身的时间复杂度是O(n),但Python在大规模数据下会暴露出语言性能的天花板。比如处理10的7次方级别的数据时,append循环的速度会比C++的同等方式慢很多。我在实际处理千万级流量日志的累计分布时,最后是切到了numpy.cumsum,速度一下子提升了十几倍:
import numpy as np arr = np.array(nums) pre_arr = np.cumsum(arr)所以我的建议是:算法刷题场景下,用纯Python的append写法就够,这样能保持代码的可读性,也避免引入numpy后在线评测系统上因库缺失或版本问题报错。但如果是本地处理真实业务数据、数据量大且有numpy环境,完全可以放开用np.cumsum。
另外,Python的整数可以无限大,计算前缀和时不会像C++那样有int溢出的问题,这算是Python一个隐形的福利,在累加大量正整数时特别省心。
6.3 刷题习惯与边界测试
我见过很多人在刷前缀和题目时,代码测了几个样例能过就急着交,结果一提交就WA。这往往不是因为思路错了,而是边界没测到。我总结了一份前缀和必测的边界清单:
- 数组长度为0或1的情况
- 查询的l等于0的情况,这最考验前缀和
pre[0] = 0的约定是否准确 - 查询区间覆盖整个数组的情况
- 数组中有大量0,或者全是负数的情况
- 差分数组操作中,区间右端点恰好等于n-1的情况,此时
r+1正好越界
每条边界都可以用几行小样例验证,别嫌麻烦。我自己的做法是写一个简单的暴力解法和前缀和解法对照着跑随机数据,两边结果不一致时再debug,效率高很多。这个方法也推荐给你,等于用一个"笨办法"来给聪明办法兜底。
7. 实战验证:四道经典题从读题到AC的完整走查
7.1 LeetCode 303:区域和检索 - 数组不可变
这道题几乎是纯前缀和的入门题。题目的要求是构建一个类,支持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:二维区域和检索 - 矩阵不可变
这是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[row2+1][col2+1] - S[row1][col2+1] - S[row2+1][col1] + S[row1][col1])我在这道题的实现上犯过一个迷糊:构建时把S[i][j]的递推关系写成了S[i-1][j-1] + matrix[i-1][j-1],漏掉了左方和上方那两个大矩形,结果查询小于3x3的矩阵时是对的,一旦矩阵变大就全乱了。所以构建二维前缀和时,一定要想清楚你算的究竟是"整个左上矩形"还是"某个小矩形",别贪图少写一个变量。
7.3 LeetCode 1109:航班预订统计(差分经典题)
这道题描述是这样的:有n个航班,用1到n编号,现在有bookings[i] = [first_i, last_i, seats_i],表示从first_i到last_i的航班每个都预定了seats_i个座位,最后返回每个航班的总预定数。
这题就是典型的"多次区间更新,最后一次性查询",直接用差分数组:
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再对应到差分数组下标。同时差分数组我开到了n+1的长度,这样即使r恰好等于n也不会越界。
这道题如果用暴力,每次预订都要遍历一遍区间,复杂度是O(n * bookings),大概率超时。而差分数组的解法只需要O(len(bookings) + n),几乎是一边遍历一边就出结果了。我强烈建议你用这道题来检验自己对差分理解的深度,因为面试里它出现的频率很高,而且换个马甲就是考同一套东西。
7.4 LeetCode 560:和为 K 的子数组(前缀和+哈希)
这道题我在第5节已经给出核心代码,这里再补一个完整的走查过程。题目问的是连续子数组中和为k的个数,注意这里子数组必须是连续的,而且元素可能有负数,所以滑动窗口直接出局。
我的思路是这样的:从左到右遍历数组,维护当前前缀和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这里最关键的一步是:先查答案,再更新哈希表。如果先更新哈希表再查,就会把当前这个位置自己也算进去,导致重复计数。我实际调试这道题时,就是因为把这两行顺序写反了,答案始终比期望值大了一圈。后来我把中间过程打出来,才意识到当前前缀和不能和它自己匹配。
这道题吃透之后,你再看一些类似的"子数组和统计"题目,会发现套路高度一致:哈希表存历史前缀和,遍历时做差查表,再更新当前前缀和。这就是一法通、万法通的效果。
我自己实际用Python跑过这四道题的完整流程,从读题到AC基本都在五分钟以内,核心的思考时间几乎全部花在"确认这道题是前缀和/差分场景,还是需要更复杂的数据结构"上。等你把前缀和的几个变体都练熟了,这种判断就会变成一种直觉——看到区间查询想前缀和,看到批量区间更新想差分,看到子数组计数想哈希表辅助,一秒钟就能定下方向。