ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

前缀和+哈希表:子数组求和与树上路径计数的O(n)解法

前缀和+哈希表:子数组求和与树上路径计数的O(n)解法 做算法题和工程性能优化的人应该没少跟“子数组求和”和“树路径求和”这两个问题较劲。我指的不是单纯调接口而是当数据规模上来之后怎么把O(n²)的暴力枚举硬生生压到O(n)或O(n log n)。这几年刷题和做线上性能优化我反复用到同一套组合拳前缀和prefix sum配合哈希表hash map。这套思路能干净利落地解决一大类“连续区间/路径上满足某种条件的计数或查询”问题而且代码量不大逻辑一旦理顺基本不会写错。这篇文章我想把这套东西彻底讲透。我会从最基础的子数组问题出发一路推到树上的路径求和与路径计数把前缀和的本质、哈希表为什么能加速、树上怎么维护以及实际调试时容易踩的坑都过一遍。不管你是在刷LeetCode准备面试还是在业务代码里做实时统计、区间查询优化这套思路应该都能派上用场。1. 前缀和不是只能求区间和从暴力枚举说起1.1 暴力解法的问题在哪里先看一个最经典的问题给一个整数数组 nums统计有多少个子数组的和等于 k。很多新手第一反应是枚举所有子数组然后算和。这当然能过小数据但一旦数组长度到了10⁵量级就彻底完蛋了。我们算一下复杂度。假设数组长度为 n子数组由左端点 L 和右端点 R 确定总共有 n(n1)/2 个。如果你对每个子数组再去一次循环求和总复杂度是 O(n³)就算你用滚动方式累计区间和也只能压到 O(n²)。n 10⁵ 的时候O(n²) 意味着 10¹⁰ 次运算在主流OJ平台上基本属于超时预定。为什么暴力会慢因为大量区间被重复计算了。比如 [L, R] 的和和 [L, R1] 的和明明只差一个元素但暴力做法里它们彼此独立没有把已经算过的结果利用起来。这就是前缀和登场的原因它把“区间查询”的耗时从 O(n) 甚至 O(n³) 直接变成 O(1)。1.2 前缀和的核心公式前缀和数组 pre 的定义很简单pre[0] 0 pre[i] pre[i-1] nums[i-1]注意我这里的写法pre[i] 表示原数组前 i 个元素的和也就是说 pre[0] 永远是 0pre[1] nums[0]pre[2] nums[0] nums[1]以此类推。有了 pre任意子数组 nums[L..R]下标从0开始的和就可以用两步算出来sum(L, R) pre[R1] - pre[L]这件事本身不难我见过很多人把它当成一个“求区间和的工具”就完事了。但真正的关键在于换个角度看这个公式如果你固定右端点 R那么“区间和等于 k”这个条件就等价于 pre[R1] - pre[L] k进一步移项得到pre[L] pre[R1] - k也就是说在遍历数组的过程中只要知道前面所有位置上有多少个前缀和的值等于“当前前缀和 - k”就能直接算出以当前位置为右端点的合法子数组数量。这个问题从“枚举两个端点”变成了“查找历史前缀和的出现次数”于是自然引出了哈希表。2. 哈希表一出手子数组计数从O(n²)变成O(n)2.1 为什么哈希表能接住这个活哈希表在这里做的事非常朴素维护“某个前缀和值出现了多少次”。遍历数组时每到一个位置我们不需要回头枚举左端点只需要查一下哈希表里有没有目标值有就累加次数。我直接给出“和为 k 的子数组个数”的完整C代码int subarraySum(vectorint nums, int k) { unordered_maplong long, int cnt; cnt[0] 1; // 前缀和为0的出现次数初始为1对应空数组 long long pre 0; int ans 0; for (int x : nums) { pre x; // 查找历史中出现过多少个前缀和等于 pre - k auto it cnt.find(pre - k); if (it ! cnt.end()) { ans it-second; } // 当前前缀和也要登记进去 cnt[pre]; } return ans; }代码短但每一行都有讲究。先说 cnt[0] 1。为什么要初始化因为当 pre - k 0 时说明从数组开头到当前位置这一段本身就是 k这时候应该被计入。如果不初始化这种子数组就被漏掉了。再说循环里的“先查后插”。比如数组只有一个元素nums [1]k 1。遍历时 pre 1先查 cnt[1 - 1] cnt[0] 1ans 变 1然后才把 cnt[1] 登记为 1。如果顺序反了先插入 cnt[1] 1再查 cnt[0] 当然也没问题但在另一些情况下会把自己算进去。举一个反例nums [1]k 0。先插入的话 pre 1cnt[1] 1查 cnt[1 - 0] cnt[1] 1结果会错误地认为有一个和为0的子数组。而正确的做法是先查因为查询时“当前前缀和”还没有登记查到的都是历史位置自然保证了子数组非空、连续且右端点在当前。2.2 手推一遍模拟过程拿 nums [1, 2, 1, -2, 3]k 2 来模拟。初始cnt {0: 1}pre 0ans 0。x 1pre 1查 cnt[1-2] cnt[-1]没找到登记 cnt[1] 1。x 2pre 3查 cnt[3-2] cnt[1]找到了cnt[1] 1ans 1。这对应子数组 [1, 2] 的和是 3等等这里 cnt[1] 对应的是第一个位置pre[1] 1pre[2] 3差值 2所以子数组是 [2]不是 [1,2]。这里容易搞混我用下标拆一下。遍历完前两个数pre 3 代表前两个元素总和是3。我们现在要找以当前元素2为右端点、和为2的子数组也就是找到一个左端点 L使得 pre[L] pre[2] - 2 1。pre[1] 1所以 L 1区间 [1,1] 就是 [2]。没错答案是1。x 1pre 4查 cnt[4-2] cnt[2]没找到登记 cnt[4] 1。x -2pre 2查 cnt[2-2] cnt[0]找到了cnt[0] 1ans 2。这对应子数组 [1, 2, 1, -2] 整个数组从开头到当前和为2。x 3pre 5查 cnt[5-2] cnt[3]没找到登记 cnt[5] 1。最终 ans 2。你可以手算验证确实只有 [2] 和 [1,2,1,-2] 两个子数组满足和为2。2.3 复杂度对比方法时间复杂度空间复杂度适用规模三重循环暴力求和O(n³)O(1)仅本地小数据前缀和 枚举左右端点O(n²)O(n)勉强到 n10⁴前缀和 哈希表O(n)O(n)n10⁵~10⁶ 稳当这套“当前状态减历史状态等于目标值”的模型能套到很多题上。比如“和为k的子矩阵”“积为k的子数组”“可被k整除的子数组个数”等核心都是把区间条件改写成一个等式再用哈希表维护历史状态的频次。3. 把前缀和搬到树上路径和的另一种打开方式3.1 从根到节点的前缀和数组的前缀和是一维的因为数组天然有线性顺序。但树不同树的路径有两个方向从上往下、从下往上。怎么把前缀和的“区间差”思想移植过来先明确一个基本概念定义 rootSum[u] 为从根节点 root 到节点 u 这条简单路径上所有节点权值之和点权和。这个值可以在一次 DFS 中算出来void dfs(int u, int parent, long long cur) { cur val[u]; rootSum[u] cur; for (int v : adj[u]) { if (v parent) continue; dfs(v, u, cur); } }如果题目给的是边权就让 cur 累加边的权值代码基本一样只是边界条件不同。有了 rootSum最直接的好处是求某个节点 u 到根节点的路径和是 O(1) 的。但更多题目问的是任意两个节点 u 和 v 之间的路径和。这时不能直接 rootSum[u] - rootSum[v]因为 u 和 v 可能分属不同子树路径不是祖先关系。3.2 任意两点路径和的公式设 L LCA(u, v)即 u 和 v 的最近公共祖先。从 u 到 v 的路径可以拆成三部分u → LL → 根根 → LL → v如果直接用 rootSum[u] rootSum[v]那么根到 L 的部分被重复算了两次所以要先减掉两次 rootSum[L]再补回 L 这个节点本身的权值因为 L 被减多了。公式如下pathSum(u, v) rootSum[u] rootSum[v] - 2 * rootSum[L] val[L]如果是边权公式会变成pathSum(u, v) rootSum[u] rootSum[v] - 2 * rootSum[L]因为边权路径不会把 L 节点本身算进去补回的那一项就不需要了。3.3 为什么 LCA 是这个方案的关键光有 rootSum 还不够要快速求任意两点路径和你必须在 O(log n) 或者更短的时间内求出 LCA。常用手段有倍增法、Tarjan离线法、树链剖分等。我平时用得最多的是倍增法预处理好 depth 数组和 fa[u][j]u 向上跳 2^j 步到的祖先单次查询 LCA 的复杂度是 O(log n)。这样整个查询路径和的流程就变成了一次 DFS 预处理 rootSum、depth、fa 数组O(n log n)。每个查询先用倍增求 LCA再用上面的公式算出 pathSumO(log n)。如果查询次数非常多可以进一步用欧拉序 RMQ 把 LCA 查询压到 O(1)于是单次路径和查询也是 O(1)。这种优化在竞赛里很常见工程上如果数据量够大也可以参考。3.4 把“路径差”和“哈希表”结合起来前面讲的是“查询固定路径和”只要一个公式就够。但如果问题是“统计满足某种条件的路径数量”只靠公式还不够还需要配合哈希表。这就进入更关键的部分了。4. 树上路径计数一维解法上树的完整实战4.1 问题重述LeetCode 437 是这类题目的典型代表给定一棵二叉树每个节点存储一个整数可正可负统计路径和等于 targetSum 的路径数目。这里的路径定义为从树上任意节点出发沿父子关系向下走到任意节点结束的一条连续路径。也就是说路径的两个端点必须是祖先-后代关系而且方向必须是从祖先到后代。如果延续“子数组和为 k”的思路我们很容易想到固定一个终点往祖先方向找有多少个起点使得从起点到终点的路径和为 targetSum。一维数组里有“前缀和出现次数”树上同样可以维护“从根到当前节点的路径前缀和出现次数”。4.2 暴力做法和它的瓶颈最容易想到的暴力是以每个节点为起点向下 DFS 寻找所有终点累加路径和等于 targetSum 的条数。对于一个根节点有 n 个节点、形状是链的树复杂度会退化到 O(n²)因为每层都要向下扫一遍。LeetCode 上节点数一多就会超时。另一种想法是固定终点为当前节点往上爬祖先节点统计 rootSum[cur] - rootSum[ancestor] targetSum 的 ancestor 数量。如果不做任何优化每个节点都要爬到根最坏也是 O(n²)。于是自然想到用哈希表保存“当前路径上所有祖先节点的 rootSum 出现次数”。4.3 前缀和 回溯的完整解法关键点在于树上的“历史前缀和”不是全局的而是当前 DFS 递归栈中从上到下的那一条路径。也就是说当你从 root 一路向下走到节点 u 时哈希表中只应该保存从根到 u 的所有祖先包括根和 u的 rootSum。一旦你从 u 回溯到父节点就必须把 u 这个 rootSum 从哈希表里删除否则当 DFS 进入兄弟分支时u 的祖先前缀和会被错误地当成“当前路径上的历史状态”。完整C代码如下class Solution { public: int ans 0; unordered_maplong long, int cnt; int pathSum(TreeNode* root, int targetSum) { cnt[0] 1; dfs(root, 0, targetSum); return ans; } void dfs(TreeNode* node, long long cur, int targetSum) { if (!node) return; cur node-val; // 找历史中是否有 rootSum cur - targetSum auto it cnt.find(cur - targetSum); if (it ! cnt.end()) { ans it-second; } // 当前 rootSum 进入路径历史 cnt[cur]; dfs(node-left, cur, targetSum); dfs(node-right, cur, targetSum); // 回溯撤销当前节点对哈希表的贡献 cnt[cur]--; } };4.4 为什么要 cnt[cur]-- 而不是直接 erase我见过很多人把回溯写成 cnt.erase(cur)在大部分测试用例下结果也是对的但这是个隐患。因为树中不同节点完全可能有相同的 rootSum 值。例如有两个祖先节点 A 和 B 的 rootSum 都等于 10哈希表里 cnt[10] 2。当你在 A 的某个子树中时哈希表里记录的是 2如果递归到某个终点后直接 erasecnt[10] 变成 0但 B 的 rootSum 还在路径上虽然 B 可能是 A 的后代或祖先关系不同这时就可能漏计数。更严谨的做法是进入节点时 cnt[cur]离开节点时 cnt[cur]--。如果减到 0保留一个值为 0 的键也不影响 find 的结果因为 find 会返回 end()。实际工程里为了避免 unordered_map 因为键太多导致哈希桶膨胀可以顺手在 cnt[cur] 0 时 erase但逻辑上先减再判断是否清除。4.5 用手推小树验证逻辑来一棵二叉树10 / \ 5 -3 / \ \ 3 2 11 / \ \ 3 -2 1targetSum 8肉眼看看有多少条。5 → 3和为 8。5 → 2 → 1和为 8。-3 → 11和为 8。所以答案是 3。用代码逻辑走一遍DFS 到节点 310的左子树的左子树的左节点这里指 rootSum18 的那个3时cur 10 5 3 18cur - 8 10哈希表里 rootSum10 出现了一次来自根节点10于是 ans 加 1对应路径 5 → 3。DFS 到节点 15的右孩子时cur 10 5 2 1 18cur - 8 10同样命中一次对应路径 5 → 2 → 1。DFS 到节点 11 时cur 10 (-3) 11 18cur - 8 10又命中一次对应路径 -3 → 11。三次命中答案正确。这个过程很清楚地展示了“当前前缀和减历史前缀和等于 targetSum”是如何在树上工作的。4.6 几个容易忽视的变形如果题目给的不是二叉树而是多叉树代码几乎不用改for 循环遍历子节点即可。如果权值在边上哈希表维护的就是“从根到父节点的路径上边权和”公式里不需要补节点的权值。如果 targetSum 可能为负数代码也完全不用动因为哈希表查的是差值不依赖数值正负。不过有一种情况要特别注意题目如果要求路径的起终点是任意两个节点不要求祖先后代关系那模型就完全不同了通常是树上点分治的范畴需要用到重心分解。可别把“向下路径”和“任意路径”混为一谈。5. 实战中容易踩的坑和我的排查思路5.1 哈希表初始化遗漏漏写 cnt[0] 1 是最经典的错误。一维数组场景下会漏掉“从开头到当前”的整段子数组树场景下会漏掉“从根到当前节点”的整条路径。这类错误的特征非常明显目标值恰好出现在数组前缀或根路径上时答案总是少 1。排查技巧拿一个两个节点的简单样例比如数组 [3, 1]k 3或者树 [1, -1]target 1。如果输出比预期少第一反应就看初始化。5.2 先查后插的顺序被颠倒前文说过先插入再查询会把“空子数组”或“空路径”也算进去。这种 bug 非常隐蔽因为只有当 k 0 或者 targetSum 0 时才会触发。我的建议是把这个顺序当成肌肉记忆先 find再插入。5.3 哈希表统计频次时误用 operator[]C 的 unordered_map 里operator[] 在键不存在时会插入一个默认值对 int 来说是 0所以如果你写成ans cnt[pre - k]; // 先访问 cnt[pre]; // 后插入第一行看起来没问题但 cnt[pre - k] 已经把不存在的键插进去了哈希表膨胀得很快。在数据量大的时候会白白损失性能严重时还会导致内存膨胀。正确写法是 find 或 count 判断或者至少像上面代码里那样用 auto it cnt.find()。5.4 整数溢出前缀和的绝对值可能很大。如果数组元素范围覆盖 ±10⁹n 10⁵前缀和能到 ±10¹⁴超出 int 范围。树路径同理。所以我所有代码都用 long long 来存 cur 或 pre哈希表的 key 也用 long long。千万别图省事用 int线上跑挂了找半天原因最后发现是溢出很冤。5.5 unordered_map 被卡常数的应对很多人以为哈希表 O(1) 是万能的但 unordered_map 的常数其实不小。在极端情况下比如 n 10⁶且前缀和没有太多重复值unordered_map 会比手写的离散化数组慢一个量级。如果笔试或项目里遇到性能瓶颈我常用的优化手段是先用一个 vector 把所有可能出现的前缀和收集起来排序后去重再用二分或数组下标映射到连续整数这样查询就是 vector 的 O(1) 下标访问。如果值域本身有限比如前缀和对某个数取模后的结果直接开定长数组。换用更轻量的哈希实现比如 C 里自定义快速哈希函数或者使用 __gnu_pbds 的 gp_hash_table。多数 LeetCode 场景不需要这些优化但工程上如果跑的是核心链路unorder_map 的扩容和时间波动确实值得警惕。5.6 树递归深度爆栈链状二叉树深度可能达到 n默认递归栈在 n 10⁵ 时就可能爆。两个办法一个是在递归函数里改为显式栈做迭代 DFS另一个是在一些允许的平台上把栈空间调大。LeetCode 通常不会卡这个但如果你在本地跑深树记得提前考虑。6. 想清楚这些扩展前缀和才算真正用活6.1 二维前缀和子矩阵和一维前缀和能算子数组和二维就是一个自然延伸用容斥原理求子矩阵和。定义 s[i][j] 表示从 (0,0) 到 (i,j) 的矩阵和那么任意子矩阵的和等于sum(x1, y1, x2, y2) s[x2][y2] - s[x1-1][y2] - s[x2][y1-1] s[x1-1][y1-1]这个公式的四个项其实就是“二维的区间差”。遇到统计满足条件的子矩阵数量的题目通常还要配合哈希表但维度升上去之后枚举方式会变得更复杂不是简单套用一维代码能解决的。6.2 差分数组动态区间修改前缀和擅长的是“静态查询”如果区间会频繁修改前缀和就力不从心这时通常要上差分数组或线段树。不过差分数组和前缀和其实是互为逆运算的对差分数组求前缀和能还原原数组对原数组求前缀和能快速查询区间和。理解这对关系你在数组题里就不容易选错工具。6.3 树上路径求和的高级扩展如果树上还同时存在点权修改前缀和就“过期”了得用树链剖分加线段树或树状数组来维护。如果只是查询也可以用树状数组做树上差分把“路径修改”转成“到根路径的修改”。这些高级玩法基础还是前缀和思想只是实现载体变了。6.4 我的核心体会回到“前缀和 哈希表”这个方法本身我个人觉得最值得记住的并不是某个模板而是那个转换思路想要求区间或路径上满足某种数量关系的子结构先把问题改写成“当前累计值 - 历史累计值 目标值”然后用数据结构去维护历史累计值的分布。一维数组的子数组、树上的向下路径、二维矩阵的子矩阵本质上都在复用同一条主线。这种思维在真实业务里也很常见。比如统计一段时间内某个指标累计达到特定增幅的次数本质就是前缀和加哈希分析用户行为序列里满足某种模式的片段也离不开类似套路。所以别看它只是一个算法小技巧掌握了它你在很多场景里都能更快地识别出问题的本质结构直接给出高效解法。
返回列表