尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

树上背包合并优化:从O(n·M²)到O(n²)的关键Trick与实现细节

树上背包合并优化:从O(n·M²)到O(n²)的关键Trick与实现细节 做树形 DP 的同学应该都有类似的经历题面一看到“在一棵树上选若干个点”或者“背包容量不超过 M”第一反应就是dp[u][j]表示 u 子树内的某个状态然后 dfs 里套两层循环去合并儿子。照着模板敲完本地一跑n5000、容量5000 直接卡死调了半天发现瓶颈全在那个三重循环上。后来听人提到一个 Trick合并儿子时把枚举上界从 M 缩到“当前已经合并过的子树大小”总复杂度就能从 O(n·M²) 降到 O(n²)。这个 Trick 表面上只是两行边界的变化背后的复杂度证明却非常漂亮而且实现时有很多细节容易记错、写错。这篇文章就把“树上背包合并”这件事彻底讲透先给结论再给证明最后把我在实战里踩过的几个坑一次性列出来。1. Trick 是什么从一个 O(n·M²) 的模板说起1.1 先统一问题模型树上背包最常见的模型长这样给定一棵 n 个点的树每个点有一个权值或者体积、价值要求在这棵树上选出一个满足某种连通性条件的点集使得某个目标函数最优。最经典的两种问法模型 A选一个大小恰好为 k 的连通块最大化点权和。模型 B每个点有体积和价值在 u 的子树内选一些点要求 u 必须选总体积不超过 M最大化总价值。这两个模型的状态定义和转移方程几乎是同一个模板dp[u][j] 表示在以 u 为根的子树中已经处理过的部分里选取了一个包含 u 且总大小为 j或总体积为 j的合法集合的最优值。“合并儿子”的操作本质上就是两个集合的合并把 u 这边已有的状态和儿子 v 子树里的状态拼在一起。这个“拼”的过程就是对左侧状态 i 和右侧状态 j 做一次组合所以转移长这样dp[u][i j] max(dp[u][i j], dp[u][i] dp[v][j]);问题在于i 和 j 各有多少个取值1.2 朴素三重循环到底慢在哪最直觉的写法是既然状态和容量 M 有关那就每次都枚举到 M。void dfs(int u, int fa) { dp[u][1] w[u]; for (int v : g[u]) { if (v fa) continue; dfs(v, u); for (int i M; i 1; --i) { for (int j 1; j M; j) { if (i j M) { dp[u][i j] max(dp[u][i j], dp[u][i] dp[v][j]); } } } } }这个是标准的 O(n·M²)。当 n5000、M5000 时总枚举量是 5000 × 25,000,000 1.25 × 10¹¹这个数量级在现代评测机上也是不可接受的。更关键的问题是绝大多数枚举都是在做无用功。因为 dp[u] 数组里真正有意义的状态只有“u 当前已经合并完的所有节点数”那么多。一个大小为 a 的集合它可能产生的状态数最多也就是 a 种不可能凭空冒出第 a1 种状态。你把 i 枚举到 M访问到的不是-INF就是永远不可能用到的无效状态。所以 Trick 的核心思想只有一句话每次合并时i 只枚举到“u 当前已经合并的子树大小”j 只枚举到“v 的子树大小”。这不是一个玄学剪枝而是一个能改变复杂度量级的正确边界。2. 核心 Trick合并时把上界改成子树大小2.1 标准写法临时数组版我用临时数组来写因为这个版本不容易写错而且后面讲循环顺序时更好理解#include bits/stdc.h using namespace std; const int N 5005; const long long NEG -(1LL 60); vectorint g[N]; long long w[N]; int sz[N]; vectorlong long dp[N]; void dfs(int u, int fa) { sz[u] 1; dp[u].assign(2, NEG); // 下标 0 和 1 dp[u][1] w[u]; // 只包含 u 自己的连通块 for (int v : g[u]) { if (v fa) continue; dfs(v, u); int old sz[u]; // u 侧当前已合并的节点数 int add sz[v]; // v 子树的节点数 vectorlong long ndp(old add 1, NEG); for (int i 1; i old; i) { if (dp[u][i] NEG) continue; for (int j 1; j add; j) { if (dp[v][j] NEG) continue; ndp[i j] max(ndp[i j], dp[u][i] dp[v][j]); } } dp[u] move(ndp); sz[u] old add; } }注意几个细节old是“处理当前儿子 v 之前”u 侧的大小不是最终大小。ndp的长度只需要old add 1不要开成M 1。dp[u]在合并完成后才更新所以dp[u][i]读取的始终是“不含 v 的旧状态”不会出现合并过程中互相污染的问题。每个无效状态都是NEG转移时跳过防止从非法状态推出“看似合法”的结果。如果题目要求容量上限 M也很简单所有地方的上界取min(size, M)同时注意i j M这个隐含条件。这个问题我在第 3.3 节单独展开。2.2 为什么边界一缩复杂度立刻降下来只看代码很难直观感受到这个优化有多猛。我们算一个具体例子。假设树是一条链节点数 n5000。朴素写法每个节点合并时都枚举 M² 次总共大约 5000 × 2500 万 1.25 × 10¹¹。用了 Trick 之后每个节点只有一个儿子合并时左侧大小始终是 1只有 u 自己右侧大小是儿子的子树大小所以第 u 层合并的枚举次数是 1 × sz[v]。把所有层加起来1 × (n-1) 1 × (n-2) ... 1 × 1 n(n-1)/2 ≈ 1.25 × 10⁷差了整整一万倍。这还只是链状树如果是星形或者随机树总枚举量同样是 n(n-1)/2 级别绝对不会超过 n²。2.3 数组版写法的循环顺序比赛时很多选手图省事直接用二维数组不搞临时 vector。这时候循环顺序就非常关键网上流传的错误模板也不少。先给正确写法long long dp[N][N]; void dfs(int u, int fa) { sz[u] 1; dp[u][1] w[u]; for (int v : g[u]) { if (v fa) continue; dfs(v, u); // i 必须倒序j 顺序无所谓 for (int i sz[u]; i 1; --i) { for (int j 1; j sz[v]; j) { dp[u][i j] max(dp[u][i j], dp[u][i] dp[v][j]); } } sz[u] sz[v]; } }为什么i一定要倒序因为在数组上原地合并时更新目标是dp[u][i j]而i j i。如果i从小到大枚举那么当i变大时dp[u][i]可能已经被这一轮合并更新过了——它已经包含了当前儿子 v 的贡献。再用它去和 v 合并一次等于把 v 的贡献算了两次答案会偏大而且是非常隐蔽的偏大。反过来i从大到小枚举时更新的是下标更大的位置这些位置不会再作为i被读取所以不会发生“新状态参与本轮合并”的问题。j的方向则无所谓因为dp[v][j]在合并过程中是只读的。提示如果你觉得倒序这个事容易记混那就老老实实用临时数组版本。代码多几行但换来的是一劳永逸的正确性。3. 复杂度证明为什么总枚举次数是 O(n²)很多博客只给结论“按子树大小合并后复杂度 O(n²)”但不说为什么。理解证明的好处是以后遇到变种题目你能自己判断复杂度是否成立。3.1 最直观的证明点对计数把一次合并看成两个集合的笛卡尔积。合并儿子 v 时左侧集合 A 是“u 当前已经合并完的所有节点”右侧集合 B 是“v 子树的所有节点”。一次合并的枚举次数就是 |A| × |B|。于是整个 dfs 的总枚举次数可以写成Σ ( |A| × |B| )这个和式从结构上看等价于“枚举树上的点对 (x, y)然后数这个点对一共被计数了多少次”。现在看任意两个不同的节点 x 和 y。它们会在哪一步被计数答案非常漂亮在 x 和 y 的最近公共祖先LCA处恰好被计数一次。理由如下。从 LCA 的角度看x 和 y 一定分别属于 LCA 的两个不同儿子子树或者其中一个是 LCA 本身。当 LCA 处理到某个儿子子树时另一个点所在的集合已经在“已合并集合 A”里了于是这一轮合并一定会把 x 和 y 配对一次。等 LCA 把所有儿子都合并完x 和 y 就永远处于同一个大集合里了。再往上的祖先合并时它们只会作为同一个集合的内部成员出现绝不会再被分到左右两侧。因此每个点对 (x, y) 对总枚举次数的贡献是 1。总点对数就是C(n, 2) n(n-1)/2所以总复杂度是 O(n²)。这个证明完全不依赖树的形态。不管是链、星、完美二叉树还是随机树上界都一样不存在“某种树形会退化成 O(n³)”的情况。3.2 把证明翻译成“合并代价之和”有的同学喜欢把复杂度写成每条边的合并代价之和总代价 Σ_{每条树边 e 对应的合并} sz(合并前 u 侧) × sz(v 侧)这个式子看着复杂其实它就是“点对数”的另一种写法。因为每个点对 (x, y) 只会对它们 LCA 的那条合并边贡献 1所以整个和式恒等于 n(n-1)/2。以后在讨论树上背包复杂度时你只需要记住只要合并时 i 的枚举范围是“u 当前已合并的大小”j 的枚举范围是“v 子树的大小”总枚举量就一定不超过 n(n-1)/2。这个结论可以直接拿来当结论用。注意这个结论成立的前提是每个点只被初始化一次状态、dp[u][j] 中 j 的范围与集合大小同步增长。如果某个写法每次都把 i 枚举到全局容量 M或者把 dp 数组的有效范围提前扩大到 n证明就直接失效了。3.3 题目给了容量上限 M 时怎么办实际题目里经常有一个额外限制总体积不超过 M。这时候状态数要截断到min(sz[u], M)。复杂度分析要分情况。情况一M ≥ n这时候截断不产生任何影响每个集合的大小就是它本身的节点数结论和前面完全一样O(n²)。情况二M 比较小每个节点保留的状态数最多 M 个。一次合并中左侧有效状态数 ≤ M右侧有效状态数 ≤ M所以一条边合并的代价最坏是 O(M²)。树有 n-1 条边所以最坏总复杂度是O(n · M²)这里非常容易踩坑很多人想当然地以为截断到 M 之后复杂度是 O(n·M)其实不是。只有在树特别“链状”的时候才接近 O(n·M)因为链状树上每次合并左侧只有一个点枚举量是 1 × min(sz[v], M)加起来是 O(n·M)。但如果树比较平衡比如一个满二叉树每个内部节点合并两个大小为 M 的子树时左右各要枚举 M 个状态一次合并就是 M²总代价就是 O(n·M²)。所以当你看到 n1e5、M500 这种数据时O(n·M²) 2.5 × 10¹⁰铁定超时。这时候要么换模型要么考虑别的优化比如依赖背包里的“返祖合并”技巧、长链剖分等不能指望这个 Trick 硬扛。4. 实操经验初始化、内存与常数优化4.1 初始化到底该用 -INF 还是 0很多人写树上背包时答案全 WA不是转移写错而是初始化的锅。这个模板的状态含义是“包含 u 的连通块”所以 u 自己必须被选。于是dp[u][1]初始化为w[u]。除了dp[u][1]其他所有状态都应该初始化为-INF表示不可达。dp[u][0]一般也是-INF因为本模板不允许出现“选 0 个点”的连通块。如果用vector版本assign(2, NEG)会把下标 0 和 1 都设成 NEG然后再手动把dp[u][1]设成w[u]。但要注意有些题目允许某个子树整体不选比如“在树上选若干个连通块最大化价值和”。这时候状态定义变成“u 的子树中可以一个都不选”那就需要额外一个dp[u][0] 0的状态并且转移时 i 要从 0 开始枚举。这属于模型变种套模板前务必先看清楚题面要求。4.2 二维数组还是 vector这是所有写树上背包的人都纠结过的问题。固定二维数组long long dp[N][N]的优点是访问快、代码简单缺点是内存巨大。n5000 时5005 × 5005 × 8字节大约是 200MB很多 OJ 直接 MLE。n3000 时也要 72MB勉强能过。vector 版本每个节点只存自己需要的那一段总内存是所有sz[u]之和乘以 8 字节。最坏情况下一条链是 n(n-1)/2 × 8 字节n5000 时大约 100MB比二维数组省一半。但 vector 的resize、assign、move都有一点常数开销一般不影响通过。我的习惯是n ≤ 3000 且内存充足直接用二维数组。n ≤ 5000 或者内存卡得紧用 vector 版本。n ≥ 1e5 还想用 O(n²) 的树上背包那基本不现实得另想办法。4.3 三个能明显提速的小细节第一跳过无效状态。在转移循环里加一行if (dp[u][i] NEG) continue;和if (dp[v][j] NEG) continue;。虽然 O(n²) 的总枚举量本来就不会访问太多无效状态但加上这个判断能避免把NEG w这种值传出去也让代码逻辑更清晰。第二临时数组只清需要的长度。如果用vectorlong long ndp(old add 1, NEG)它只清old add 1个元素。千万不要在循环里动不动就memset(dp[u], 0, sizeof(dp[u]))那会把整行 N 个元素都清一遍把一个 O(n²) 算法搞成 O(n² n·N) 甚至更糟。第三sz[u] sz[v]的位置一定放在合并完成之后。如果你在合并之前就累加sz[u]那么循环里i的上界会超过 dp[u] 的有效长度访问到一堆未初始化的状态轻则复杂度退化重则答案错乱。5. 常见问题与排查思路速查表5.1 症状对照表症状可能原因解决方向大数据 TLE合并循环的上界写成了 M 而不是 sz检查所有for循环边界确认是否用了sz[u]、sz[v]答案偏大且小数据正常原地数组合并时 i 用了正序把 i 改为倒序或改用临时数组版本答案全是 NEG 或极端负值dp[u][1] 没初始化或初始值设成了 0检查初始化逻辑确认不可达状态是 NEG总枚举次数超过 n² 很多合并前就更新了 sz[u]或者临时数组开成了 N 大小检查sz[u] sz[v]的位置MLE二维数组 N×N 超内存改用 vector 按需分配检查 N 是否过大答案比理论最优小dp 初始值用的 0非法状态参与了转移把非法状态设为 NEG并在转移前跳过5.2 一个真实的排查过程有一次我做一道 n5000 的树上连通块 DP本地测试数据全对一交上去 TLE。当时我第一反应是常数太大于是给合并循环加了计数器long long cnt 0; // 在两层循环内部 cnt;跑完之后打印cnt发现是 2.5 × 10⁸而不是预期的 1.25 × 10⁷。我马上意识到上界写错了。回去一看代码里写的是for (int i sz[u]; i 1; --i) for (int j 1; j sz[v]; j)看起来没问题但问题出在sz[u]在合并前就已经被更新过了。原来我为了省事在 dfs 开头就先处理完所有儿子的sz[u]合并时sz[u]已经是整棵子树的完整大小于是每个节点的合并都多枚举了很多。改成“边合并边累加sz”之后计数器立刻变成了 1.25 × 10⁷ 级别。这个坑很隐蔽因为结果是对的只是慢。以后遇到树上背包超时第一件事不是优化常数而是先写个计数器看总枚举次数是否符合 O(n²)。6. 这个 Trick 的扩展用法与我的习惯6.1 树上依赖背包一样适用如果题目要求“选一个点之前必须先选它的父节点”比如经典的选课问题状态会变成dp[u][j]表示“在 u 的子树中选 j 个点且 u 必须被选”。转移时同样按子树大小合并复杂度同样是 O(n²)。区别只在初始化依赖背包里dp[u][1]必须先赋值为 u 的权值转移时 i 必须从 1 开始因为 u 必须选不能出现dp[u][0]这种状态。如果你把 0 状态也加进去就可能在合并时构造出“没选父节点却选了子节点”的非法方案。6.2 其他树形 DP 优化与本 Trick 的边界O(n²) 的树上背包合并适合 n ≤ 5000 左右的题。如果 n 到 1e5就得考虑更高级的优化长链剖分把与深度相关的 DP 优化到 O(n)适用于某些特定转移。dsu on tree处理静态子树查询常用于计数类问题。依赖背包的“重儿子优先”写法常数上比普通写法小一些但复杂度量级不变。不要指望一个 Trick 通吃所有树形 DP。这个 Trick 的价值在于它把“树上背包合并”的最坏复杂度从 O(n·M²) 锁定到了 O(n²)而且证明简洁、实现成本低。它是树形 DP 选手的必修基本功但绝不是终点。最后分享一个我自己写树上背包时的小习惯不管题目有没有给 M我都会在 dfs 里维护sz[u]把循环上界写成sz的表达式。这样做不仅是为了复杂度正确更是强制自己时刻清楚“当前状态里到底包含了多少个节点”。状态定义清晰了转移自然不容易写错。如果你刚开始学树上背包建议这个 Trick 不要只背模板而是要像第 3 节那样亲手把点对计数的证明推一遍。想明白“每个点对恰好被枚举一次”之后以后再遇到任何变种你都不会再被复杂度问题困住了。
返回列表