ARTICLE DETAIL

资讯详情

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

换根DP破解力扣3772:从O(n²)超时到O(n)高效解法

换根DP破解力扣3772:从O(n²)超时到O(n)高效解法 做图论题最怕的不是算法有多难而是你一眼看出了正确方向却把复杂度写没了。力扣 3772 这道“子图的最大得分”标题里的图论、DFS、换根法三个词基本就把解法写在脸上了但真正动手实现的时候仍然有不少细节会让人翻车。我第一次提交的版本非常朴素枚举每个节点当根每个根都跑一次树形 DP统计包含这个根的最大连通子图得分然后取最大值。样例秒过一提交(O(n^2)) 直接超时。n 到 (2\times 10^5) 的规模这个复杂度等于当场枪毙。正解就是换根法也就是常说的二遍 DFS / Rerooting DP先固定一个根做一次自底向上的树形 DP再做一次自顶向下的换根遍历n 个根各自的答案在一次遍历里全部算完。这篇文章我会从朴素做法开始逐步推出换根公式给出 Python 和 C 的完整代码最后重点讲几个刷题时最容易踩的坑。1. 先看懂题目考什么1.1 题目模型的一个关键抽象先说我理解并使用的题意抽象。给定一棵 n 个节点的树节点编号 1 到 n每个节点 i 有一个整数权值 (a[i])权值可以是负数。对于任意一个节点 r我们选择树上的一个连通点集 S要求 r 必须包含在 S 中。子图的“基础得分”是 S 内所有节点的权值之和。为了让“根”这个身份真正影响答案题面设定根节点 r 的权值在计算总分时要额外再加一次也就是说总得分等于[ score(S, r)\sum_{v\in S}a[v]a[r] ]其中 (r\in S)。如果你的版本里没有“根节点额外计一次分”这个条件那也没关系核心换根流程完全一样只需要把最后输出答案时的 a[u]去掉即可。后面先按这个带额外条件的版本讲因为它能天然解释“为什么要枚举根”根不同答案不同哪怕 S 完全一样最后的得分也可能因为根权值多算一次而发生变化。如果只是求一个全局最优的任意连通子图一次普通的树形 DP 就够了根本不需要换根。但当我们要求“每个节点作为根时的最优值”或者题目答案本身和根挂钩时就必须高效地求出所有根的答案换根法就是为这种情况设计的。1.2 朴素做法为什么必死朴素思路非常直接对每一个节点 r都从 r 出发做一次 DFS把以 r 为根时包含 r 的最大连通块权值和算出来。伪代码大概是这样的for r in 1..n: ans max(ans, dfs(r, -1))其中dfs(u, fa)返回以 u 为“当前根”、必须包含 u 的最大连通块分def dfs(u, fa): cur a[u] for v in g[u]: if v fa: continue t dfs(v, u) if t 0: cur t return cur这个朴素版正确性没问题。任何一个包含 u 的连通块都可以看成“u 本身”加上它在某些邻居方向上的连通子块各邻居之间没有边相连只能通过 u 连通所以决策是独立的。分支收益为正就收为负就丢。但问题是复杂度。每次都遍历整棵树每次遍历要访问 (2(n-1)) 条无向边总复杂度 (O(n^2))。当 (n2\times 10^5) 时操作量是 (4\times 10^{10}) 量级无论如何都跑不完。我一开始没有意识到问题的严重性还想着“反正样例小应该没事”结果一发提交直接教做人。重复计算也特别明显比如从根 1 换到根 2节点 3 的整棵子树完全没变但朴素做法把节点 3 下面所有节点又重算了一遍。换根法解决的就是这个浪费把和“变化的分支”无关的那部分结果缓存下来只更新真正发生变化的信息。2. 固定根的一次 DFS把一棵子树算透彻2.1 状态定义与转移方程既然要换根第一步永远是把一个固定根下的信息算清楚。这里固定根选 1。定义数组down[u]在原树以 1 为根的结构中以 u 为根的子树里必须包含 u 的最大连通块权值和。注意“必须包含 u”这六个字它保证了我们可以用子树答案拼出父节点的答案。转移方程是[ down[u] a[u] \sum_{v\in child(u)} \max(0, down[v]) ]其中 child(u) 是在以 1 为根时 u 的邻居中那些不是父节点的节点。理解这个方程需要抓住一个点u 的每个孩子 v 所在的区域是互相独立的它们之间没有边只能通过 u 相连。因此 u 的全局最优连通块一定是“u 本身”加上若干“孩子方向上的最优连通块”。如果一个孩子分支自己的最优值是 -3把它接进来会让父节点总收益下降 3肯定不干如果孩子分支最优值是 4接进来就是纯赚。这也是max(0, down[v])的来源。举个我调试时用的小例子。树结构如下括号里是权值1(-2) / \ 2(5) 3(3) / \ / \ 4(-1) 5(2) 6(-4) 7(6)以 1 为根先算底部的叶子(down[4]-1)因为节点 4 没有孩子最多只能选自己。(down[5]2)。(down[6]-4)。(down[7]6)。然后算节点 2[ down[2]5\max(0,-1)\max(0,2)5027 ]节点 3[ down[3]3\max(0,-4)\max(0,6)3069 ]最后节点 1[ down[1]-2\max(0,7)\max(0,9)14 ]这里要注意(down[2]7) 表示在“1 是根”的结构里从 2 出发并限制在 2 的子树内最多能拿 7 分对应的点集是 {2, 5}。同理 (down[3]9) 对应 {3, 7}。而节点 1 的 14 分对应 {1, 2, 3, 5, 7}权值和是 (-2532614)。负数节点 4 和 6 都没有被选进来。2.2 为什么分支收益要用 max(0, ·)max(0, down[v])这个写法是整道题的一半核心也是很多初学者最容易想当然写错的地方。如果把转移写成[ down[u] a[u] \sum_{v\in child(u)} down[v] ]那就默认了每个孩子方向都必须被选进来。这在全是正权值的时候没问题但权值一出现负数这行代码就会把负贡献强行加到答案里。题目里节点权值可以是负数所以这种写法从一开始就是错的。另一个容易犯的错误是写成[ down[u]\max(a[u], a[u]\sum_{v\in child(u)} down[v]) ]这也是错的因为你在做决策时不是“某个孩子整体要不要”而是“每个孩子方向分别要不要”。如果三个孩子方向的收益分别是 -5、2、3正确结果是 (a[u]023)而上面那个错误写法会变成要么只选 a[u]要么把 -5 也一起吃掉丢了“丢弃负分支”的能力。记住口诀树上的连通块求和问题永远是一个一个分支单独考虑谁正选谁谁负扔谁不要整体打包。2.3 第一次 DFS 的两种实现先给递归写法清晰直白def dfs1(u, fa): cur a[u] for v in g[u]: if v fa: continue t dfs1(v, u) if t 0: cur t down[u] cur return cur调用dfs1(1, -1)就能得到所有down。但递归有个隐患Python 默认递归深度只有大约 1000遇到 (n2\times 10^5) 的链状树会直接RecursionError。我更喜欢用显式的“栈序”来写这样可以彻底避开递归深度问题比赛和刷题环境里都更稳。做法是先做一次 BFS/DFS 遍历记录父节点数组parent同时得到一个“先父后子”的顺序order。然后倒序遍历order算down。parent [0] * (n 1) order [] stack [1] parent[1] -1 while stack: u stack.pop() order.append(u) for v in g[u]: if v parent[u]: continue parent[v] u stack.append(v) for u in reversed(order): cur a[u] for v in g[u]: if v parent[u]: continue if down[v] 0: cur down[v] down[u] cur这里order里每个节点一定排在它的所有子孙节点之前所以倒序遍历时处理 u 的时候它的孩子节点down都已经算好了。这个写法对后续换根同样适用因为我们还要正序遍历order来做自顶向下的转移。3. 换根公式把树根从父亲搬到儿子3.1 补集思想与 side 值down算完后所有“以 1 为根时的子树信息”都在手上了。现在定义[ all[u]\text{以 u 为根时包含 u 的最大连通块权值和} ]all[1]很好求因为 1 本来就是最开始选定的根[ all[1]down[1] ]问题是从 1 出发怎么把all传给儿子。比如现在要把根从 fa 换到 u。当 u 成为新根后u 面对的邻居分两类u 原来的孩子节点。这些子树结构完全没变对应的候选贡献还是 (down[v])。fa 方向。现在 fa 变成了 u 的一个“儿子分支”但这个分支不是 fa 的原始孩子分支它是“fa 除了 u 方向以外的一切”。所以关键是算出 fa 方向上那个候选值。设它为side。它等于[ side all[fa] - \max(0, down[u]) ]为什么这么减因为all[fa]是在“fa 是整棵树根”时包含 fa 的最优连通块权值和。这个最优连通块如果经过 u 方向向外扩展一定包含了 u 方向上那个最优分支。而 fa 在计算时只有在 (down[u]0) 的情况下才会真的把 u 方向的收益算进去如果 (down[u]\le 0)它根本就没选这个方向贡献是 0。所以从all[fa]中剔除 u 方向时要剔除的是 (\max(0, down[u]))而不是 (down[u])。这是换根公式里最精妙也是最容易出错的一步后面第 5 节会专门讲。3.2 公式推导与完整示例fa 方向的候选分支值算出来后u 是否接纳它就变成一个独立决策收益为正就收收益为负就丢。于是[ all[u] down[u] \max(0, side) ]代入side[ all[u] down[u] \max\left(0,\ all[fa]-\max(0, down[u])\right) ]这就是完整的换根公式。它的形式非常简洁新根的答案 自己原来的子树答案 父亲的整树答案中剔除自己后的正收益部分。继续用之前的例子验证。我们已经算到[ down[2]7,\ down[3]9,\ down[4]-1,\ down[5]2,\ down[6]-4,\ down[7]6,\ down[1]14 ](all[1]down[1]14)。换根到 2[ all[2]down[2]\max(0,\ all[1]-\max(0,down[2]))7\max(0,14-7)14 ]含义是2 成为根后原本 1 方向剩下的 {1, 3, 7} 收益 (-2367) 值得并入所以总答案 14。换根到 3[ all[3]9\max(0,14-9)14 ]换根到 4叶子节点(down[4]-1)[ all[4]-1\max(0,\ all[2]-\max(0,-1))-11413 ]这里能看到减法的妙处因为 (down[4]) 是负数原来 all[2] 根本没有把节点 4 方向算进去所以剔除的贡献是 0父亲方向整整 14 分全部保留节点 4 接住后总得分变成 13。再算节点 7[ all[7]6\max(0,\ all[3]-\max(0,6))6\max(0,8)14 ]最终每个根的全部答案都出来了。如果题目要求带“根节点额外加一次”那么答案需要取[ ans\max_{u} (all[u]a[u]) ]对照刚才的例子(u7)(all[7]14)再加 (a[7]6)得到 20。(u2)(all[2]14)再加 5得到 19。(u3)(all[3]14)再加 3得到 17。所以最佳根是节点 7答案是 20。大家可以手算一下以 7 为根选择连通块 {1, 2, 3, 5, 7}基础得分 14根节点 7 额外加 6总分 20和换根算出来完全一致。3.3 第二次 DFS 实现第二次遍历要从根 1 自上而下地执行换根公式。原因是计算all[u]依赖all[fa]而 fa 是 u 的父亲所以必须先算好父亲再算儿子。order本身就是先父后子的顺序直接正序遍历即可。递归写法def dfs2(u, fa): for v in g[u]: if v fa: continue side all[u] - max(0, down[v]) all[v] down[v] max(0, side) dfs2(v, u)对应迭代写法不需要递归函数all[1] down[1] for u in order: for v in g[u]: if v parent[u]: continue side all[u] - max(0, down[v]) all[v] down[v] max(0, side)这样两次 BFS/DFS 风格的遍历就搞定了所有根的状态核心代码不超过 20 行。4. 完整代码Python、C 和复杂度对比4.1 Python 完整实现下面代码采用的是迭代写法不依赖递归深度也顺便避免了 Python 递归爆栈的问题。import sys from collections import deque def solve(): input sys.stdin.readline n int(input()) a [0] list(map(int, input().split())) g [[] for _ in range(n 1)] for _ in range(n - 1): u, v map(int, input().split()) g[u].append(v) g[v].append(u) parent [0] * (n 1) order [] q deque([1]) parent[1] -1 while q: u q.popleft() order.append(u) for v in g[u]: if v parent[u]: continue parent[v] u q.append(v) down [0] * (n 1) for u in reversed(order): cur a[u] for v in g[u]: if v parent[u]: continue if down[v] 0: cur down[v] down[u] cur all_val [0] * (n 1) all_val[1] down[1] for u in order: for v in g[u]: if v parent[u]: continue side all_val[u] - max(0, down[v]) all_val[v] down[v] max(0, side) ans max(all_val[i] a[i] for i in range(1, n 1)) print(ans) if __name__ __main__: solve()一点说明第一步 BFS 里order中父亲一定排在孩子前面这是由 BFS 的层次顺序保证的。第二次遍历正序扫order处理到 u 时它的父亲已经处理完all_val[u]必然已经算好再算孩子就是安全的。4.2 C 完整实现#include bits/stdc.h using namespace std; int main() { ios::sync_with_stdio(false); cin.tie(nullptr); int n; cin n; vectorlong long a(n 1); for (int i 1; i n; i) cin a[i]; vectorvectorint g(n 1); for (int i 0; i n - 1; i) { int u, v; cin u v; g[u].push_back(v); g[v].push_back(u); } vectorint parent(n 1, 0), order; order.reserve(n); queueint q; q.push(1); parent[1] -1; while (!q.empty()) { int u q.front(); q.pop(); order.push_back(u); for (int v : g[u]) { if (v parent[u]) continue; parent[v] u; q.push(v); } } vectorlong long down(n 1, 0); for (int i n - 1; i 0; --i) { int u order[i]; long long cur a[u]; for (int v : g[u]) { if (v parent[u]) continue; if (down[v] 0) cur down[v]; } down[u] cur; } vectorlong long all_val(n 1, 0); all_val[1] down[1]; for (int u : order) { for (int v : g[u]) { if (v parent[u]) continue; long long side all_val[u] - max(0LL, down[v]); all_val[v] down[v] max(0LL, side); } } long long ans LLONG_MIN; for (int u 1; u n; u) { ans max(ans, all_val[u] a[u]); } cout ans \n; return 0; }C 版本里所有累加量都用long long因为权值是负数或正数叠加时完全可能超过 int 范围尤其是 n 很大时。这一点在刷题时很常见别为了省一点空间写出隐患。4.3 复杂度与朴素方案对比方案时间复杂度空间复杂度能不能过朴素枚举根 DFS(O(n^2))(O(n))(n2\times10^5) 时必超时固定根树形 DP 换根(O(n))(O(n))可以过且常数很小具体到操作次数换根法一共两次遍历每次遍历每条无向边恰好被看两次所以总操作量是 (4n) 级别的常数。以 (n2\times 10^5) 计算大概 (8\times 10^5) 次边访问运行时间在毫秒量级。对比朴素法的 (4\times 10^{10}) 级别访问差距一目了然。5. 最容易翻车的四个细节5.1 全为负权值时的答案如果所有节点权值都是负数比如 (a[-3,-5,-2,-4])那么任何包含两个以上节点的连通块只会让得分越来越差。最优策略就是单独选一个“亏得最少”的节点。在 DP 里这会自然体现为所有孩子的down[v]都小于 0max(0, down[v])全部是 0所以每个down[u]都等于自己的a[u]。换根后all[u]也等于a[u]。如果题目要求根节点额外加一次最终答案就是[ \max_{u} 2\cdot a[u] ]也就是选最接近 0 的那个负数权值节点当根。我见过不少人在这里翻车把转移写成直接把down[v]加进去结果全负数时把整棵树都选上了得到一个很大的负数和正确答案差得十万八千里。检查自己代码有没有这个问题最简单的方法就是造一组全负数的小数据跑一遍对比手算答案。5.2 叶子节点的换根过程叶子节点在换根时最容易让人困惑因为它的down就是a[u]可能是个负数。以节点 4 为例[ down[4]-1 ]换根到 4 时[ sideall[2]-\max(0,-1)all[2]-014 ][ all[4]-1\max(0,14)13 ]这里一定要理解为什么是减 0 而不是减 -1。all[2]在计算时根本没有把节点 4 这个分支选进去因为它权值为负。既然本来就没选剔除掉的东西就应该是 0。如果你写成all[fa] - down[u]就会变成 (14-(-1)15)相当于凭空多捏造出 1 分结果完全错乱。类似的换根到叶子节点时叶子节点的答案理论上只会包含两种可能要么只选叶子自己要么把父亲方向的整块收益全拿过来。因为叶子没有别的分支down[u]a[u]公式自然做到这一点。5.3 递归爆栈与迭代写法这道题 n 最大到 (2\times10^5)如果遇到一条链状的树递归深度就是 (2\times10^5)。Python 默认递归深度只有 1000直接RecursionError有些 C 环境下系统栈也不一定扛得住这么深的递归。虽然可以写sys.setrecursionlimit(10**6)但如果一棵树真的退化成链递归函数反复压栈仍然存在风险而且在某些 OJ 上 C 深递归也可能异常退出。我的建议是一开始就写成迭代形式。做法就是上面代码里的两步用队列或栈生成parent和order。倒序扫order做自底向上 DP正序扫order做换根。这样一来代码完全不含递归时间和空间都可控。刷题时省心很多不用每次提心吊胆地估算递归深度。5.4 换根减法中的 max(0, ·) 陷阱这是换根法里最隐蔽的坑。公式本身是side all_val[fa] - max(0, down[u])但有些人会不自觉地写成side all_val[fa] - down[u]在大部分情况下如果down[u]是正数两者结果一致但一旦down[u]为负差别就出来了# 正确 side all_val[fa] - max(0, down[u]) # 错误 side all_val[fa] - down[u]当 (down[u]0) 时错误的写法相当于“把一个负数分支从父答案里硬扣掉”结果会让side比真实值大。比如某个父亲方向的真实收益是 14孩子方向down[u]-1正确剔除后还是 14错误剔除后变成 15然后孩子答案就虚高 1 分甚至可能直接改变最终最大值。这种错误在小数据上不一定能测出来因为输出答案可能只是差一点但一旦数据规模大、权值分布复杂结果就会完全不对。我自己调试时曾卡在这个位置将近半小时最后就是一组随机对拍数据暴露了问题。排查方法也简单在换根循环里打印u、down[u]、all_val[fa]、side和朴素 DFS 的结果逐项对比哪里不一致一目了然。6. 换根法能解决的远远不止这一题6.1 换根 DP 的通用形态换根法本质上是一个处理“树上所有点为根”的问题模板。绝大多数题目的套路分成两步固定任意一个根自底向上 DP算出每个节点只在子树范围内的答案。自顶向下做第二次遍历把“父亲方向”的信息通过补集方式合并进来算出每个节点作为全局根时的答案。不同题目的状态定义和转移公式差别很大但骨架永远是“子树答案 父亲侧补集”。这个思想在力扣 834“树中距离之和”里是类似的在 Codeforces 1092F“Tree with Maximum Cost”里也几乎一样只是把状态从“连通块权值和”换成了“加权距离和”。换个形式状态变成计数、最小值、最大值、合法方案数换根的思路都通用。6.2 值得做的几道练习题如果你想把这个技巧彻底练熟我推荐按下面的顺序做几个题每个题都用“先固定根再换根”的思维去分析状态力扣 834树中距离之和。最经典的换根入门题状态是“子树内所有节点到当前节点的距离和”。换根公式也很好推适合练手。力扣 2581统计可能的树根数目。换根时还要维护一个“计数合法性”条件不是简单取 max但思考路径一样。力扣 2858可以到达每一个节点的最少边反转次数。每个点作为根时需要的翻转次数用换根法可以 O(n) 算出全部答案。Codeforces 1092FTree with Maximum Cost。带权距离和的换根题和 834 很像但权值不再是 1推导时注意前缀和式的转移。每道题做完后建议用朴素 DFS 写一个对拍脚本随机生成小规模树比较换根版和朴素版的输出是否完全一致。对拍是刷算法题性价比最高的自查手段特别是换根这种容易在两个max和减法符号上出错的问题。6.3 我的刷题和自测方法最后分享一下我个人的习惯。遇到换根题我不会先写完整代码而是先在草稿纸上固定一个根把down和all的含义写清楚再手算一组 5 到 7 个节点的数据。手算的过程看着笨但能帮你确认公式没推错。然后写一个朴素版作为“标准答案”再写换根版用随机数据对拍。对拍脚本本身很便宜几十行就能跑起来。数据生成时注意覆盖几种特殊情况全是正权、全是负权、单点树、链状树、完全二叉树。我踩过的几次坑几乎都是靠这些特殊数据抓出来的。尤其是“全是负数”和“链状结构”这两类基本每次都能打出一个隐藏问题。实际编码中还有一个经验把down和all都定义成long long或 Python 的 int不要在负数转移上省类型。这道题所有权值累加时int 溢出可能会把正确答案变成完全不相干的数值而且这种错误极难定位。先保证类型安全再去抠常数和优化是我写树形 DP 的一贯顺序。
返回列表