从医院设置到带权重心：树形DP换根法全解析

做算法竞赛刷题的人大概率都在洛谷上遇过P1364这道“医院设置”。题面不复杂给一棵二叉树每个节点住着若干居民挑一个节点建医院让所有人到医院的路程总和最小。这道题我刷过两遍第一遍用Floyd莽过去的第二遍才真正把树形DP和换根的思想吃透。说实话n≤100的约束让这题几乎什么写法都能AC但“能过”和“理解”是两码事——它背后藏的是带权树重心模型恰好能把最短路、BFS、树形DP三条知识线串在一起。这篇文章就把这条线完整捋一遍适合正在备战算法竞赛、刚啃完二叉树基础、想在树上DP上找一个经典切口的人。1. 题目在问什么把“医院选址”翻译成数学模型1.1 从题面到二叉树三个关键信息的提取先别急着写代码把题面的信息拆开。P1364的输入结构非常固定第一行是节点数n接下来n行每行给三个整数分别是当前节点的居民数、左孩子编号、右孩子编号孩子编号为0表示没有。这三个信息对应三种图论元素节点上的居民数是点权左孩子、右孩子构成的边是树边题目明确每条边距离是1要建的医院可以放在任意一个节点上包括居民数为0的节点。很多人第一次做这题盯着“二叉树”三个字就开始递归建树然后从根往下搜索算完就往右孩子、左孩子递归最后提交发现样例都过不了。问题就出在他把“医院只能往子树方向走”了实际上人可以从孩子走到父亲也可以从父亲走到孩子的孩子整棵树是连通的所有节点之间都可达。所以第一步建模结论是这不是一棵“有根二叉树”的问题而是一棵无向树上求带权最优点的问题。这个认知转换是整道题的分水岭。1.2 目标函数的数学意义与“距离”的陷阱如果医院设在点v总代价的数学表达是S(v) Σ( w[u] × dist(u, v) )其中w[u]是u节点的居民数dist(u,v)是u到v经过的边数。注意是“边数”不是“经过的节点数”。这个细节虽然小但写DFS累加时特别容易多算或少算一个单位。从u到v如果经过三条边距离就是3居民数乘以3这就是u点对总费用的贡献。来看个具体例子。假设一棵链状的树1号点有10个人2号点有5个人1到2是连着的。医院建在2号点总费用就是10×1 5×0 10如果医院建在1号点费用是10×0 5×1 5。同样是两个节点换个位置就差了5。这说明居民权重和距离是相乘的关系而不是简单比较两边人数。另一个陷阱是居民数为0的节点也能当医院。比如一个三节点链两端的点各住100人中间点没人那最优解一定在中间点。如果写代码时把“没人的点”跳过了反而可能错过正确答案。1.3 n≤100这个约束到底在暗示什么n≤100是故意的。这个范围意味着O(n³)的Floyd全源最短路100³ 100万次运算完全能过O(n²)的枚举BFS每个点跑一次全树BFS也完全能过O(n)的换根DP更是随便过。所以这题本质上是一道“怎么都能AC但你得知道自己在做什么”的题目。我的建议是先写一个最不容易出错的暴力版本拿分再用正解对拍验证而不是一上来就写DP写完了都不知道对不对。数据范围小的题有一个好处——非常适合用来验证思路。等哪一天碰到n10⁵的带权重心题就知道当时在P1364上把原理啃透是多值得的一件事。2. 先别急着优化三种拿分写法的复杂度与代码量2.1 写法一枚举医院BFS跑全树O(n²)这是最符合直觉的写法把每个点都假设成医院用BFS从医院出发把整棵树搜一遍边搜边累加“当前点的居民数 × 到医院的距离”。但这里有个绕不开的建模问题题目只给了左右孩子BFS怎么从某个节点出发走到它的父节点答案是读入的时候就别按“二叉树结构”存而是存成无向邻接表。#include bits/stdc.h using namespace std; const int N 105; int w[N]; vectorint g[N]; int bfs(int s) { vectorint dist(N, -1); queueint q; dist[s] 0; q.push(s); int res 0; while (!q.empty()) { int u q.front(); q.pop(); for (int v : g[u]) { if (dist[v] ! -1) continue; dist[v] dist[u] 1; res w[v] * dist[v]; q.push(v); } } return res; } int main() { int n; cin n; for (int i 1; i n; i) { int l, r; cin w[i] l r; if (l) { g[i].push_back(l); g[l].push_back(i); } if (r) { g[i].push_back(r); g[r].push_back(i); } } int ans INT_MAX; for (int i 1; i n; i) ans min(ans, bfs(i)); cout ans endl; return 0; }这段代码有一个容易被忽略的地方dist[s] 0之后进入BFS主体之前医院所在点s自己的居民数贡献是0所以BFS里没有额外加w[s] * 0这是对的。然后从s开始逐层扩展每个新节点的距离是上一个节点距离1贡献是w[v] * dist[v]。复杂度是n次BFS每次O(n)总共O(n²)n100时大约一万次边访问飞快。2.2 写法二Floyd全源最短路 枚举O(n³)如果你觉得BFS还要写队列、判重不够省心那Floyd就是最无脑的方案。思路分三步初始化dist[i][i] 0其他为无穷大左右孩子存在时把dist[i][l] dist[l][i] 1右孩子同理三重循环跑Floyd得到任意两点间的距离枚举所有节点当医院算Σ(w[i] * dist[i][j])取最小值。#include bits/stdc.h using namespace std; const int N 105; const int INF 0x3f3f3f3f; int w[N], d[N][N]; int main() { int n; cin n; memset(d, 0x3f, sizeof(d)); for (int i 1; i n; i) d[i][i] 0; for (int i 1; i n; i) { int l, r; cin w[i] l r; if (l) d[i][l] d[l][i] 1; if (r) d[i][r] d[r][i] 1; } for (int k 1; k n; k) for (int i 1; i n; i) for (int j 1; j n; j) if (d[i][k] d[k][j] d[i][j]) d[i][j] d[i][k] d[k][j]; int ans INF; for (int j 1; j n; j) { int sum 0; for (int i 1; i n; i) sum w[i] * d[i][j]; ans min(ans, sum); } cout ans endl; return 0; }Floyd写起来最容易但它的启发性最低——它把树当成普通图来处理完全没有利用树的祖先关系。不过作为对拍用例的基准实现它是称职的。2.3 三种暴力的选择逻辑这里把三种方法摆在一起看写法复杂度代码量易错点枚举BFSO(n²)中等dist重置、邻接表建双向边FloydO(n³)最少INF初始化、自环置0递归DFS暴力O(n²)中等漏掉父方向结果偏小我在比赛里如果碰到带权重心题第一反应是写BFS枚举保底然后把它当成对拍器。Floyd适合时间紧又不想动脑的时候用来“垫分”但不要把它当成学习目标——会用Floyd不算懂这题能把换根DP写出来才算真懂。3. 正解的推导树上的带权重心与二次扫描换根3.1 为什么这是一道换根DP题暴力写法把每个点都当成医院算一遍重复计算了大量信息。换根DP的核心思想是只完整计算一个点当医院的答案其余点由它的邻居答案“递推”出来。想象医院从父节点u移动到它的孩子v。这棵树的边u-v被“跨过”一次而这棵树的节点可以被分成两块v这棵子树里的所有居民到医院的路径都少走了一条边总代价减少sz[v]v子树之外的所有居民到医院的路径都多走了一条边总代价增加tot - sz[v]。这里的sz[v]是v子树内的居民总数tot是整棵树的居民总数。于是就有了全题最核心的转移公式dp[v] dp[u] - sz[v] (tot - sz[v]) dp[u] tot - 2 × sz[v]理解这个公式的等价说法是把每一条边看作“管道”医院位置移动时经过这条边的流量方向发生反转代价差就是两边人数的差。3.2 第一次遍历以1为根统计子树人口与内部代价换根DP需要两次DFS。第一次DFS以任意节点为根代码里取1号点从它向下递归求出两个东西sz[u]以u为根的子树内居民总数dp[u]仅考虑u子树内的居民让这些居民走到u的总代价。第一次DFS的转移是sz[u] w[u]; dp[u] 0; if (lc) { dfs1(lc); sz[u] sz[lc]; dp[u] dp[lc] sz[lc]; // 孩子子树内部先走到孩子再统一走一条边到u }dp[lc] sz[lc]的意思是孩子子树里每个居民要先汇到孩子节点这部分代价是dp[lc]然后再从孩子节点走一条边到u因为lc子树里有sz[lc]个居民所以额外加sz[lc]。注意这里边的代价是1所以乘以的是1不是距离层级。如果边权是c就得乘c这个后面扩展题会用到。第一次DFS跑完dp[1]就是“医院设在1号点时全体居民到1号点的总距离”。这是唯一的“全量计算”后面所有点的答案都由它推出。3.3 第二次扫描从父答案推出子答案第二次DFS从根开始利用转移公式把父节点的答案推给孩子节点。void dfs2(int u) { ans min(ans, dp[u]); if (lc) { dp[lc] dp[u] tot - 2 * sz[lc]; dfs2(lc); } if (rc) { dp[rc] dp[u] tot - 2 * sz[rc]; dfs2(rc); } }注意公式里的tot是全局居民总数就是sz[1]。每次往孩子走用父节点的dp和两个孩子各自的子树大小算出孩子节点的dp然后继续递归。整个过程只做一次全树遍历所以复杂度是O(n)。这里的tot - 2*sz[child]可以看作“换根带来的增量”。如果孩子子树的人口超过总人口的一半增量是负数说明往这个孩子方向走能让总代价变小如果不超过一半增量非负说明这个孩子方向不是更优的方向。这就是带权重心的判断条件的雏形。3.4 直观理解与验证拿一个简单链状树验证这个公式比如三个节点排成一条链1号节点300人2号节点一人没有3号节点100人左右孩子关系为1的左孩子22的左孩子3。第一次DFS后sz[3] 100dp[3] 0sz[2] 100dp[2] dp[3] sz[3] 1003号点100人走到2号每人1条边共100sz[1] 400dp[1] dp[2] sz[2] 100 100 2002号和3号的居民合计100人再加1号点300人自己实际总代价是200因为医院在1号3号100人走2条边达200。然后换根到2号dp[2] dp[1] tot - 2*sz[2] 200 400 - 200 400不对重新算一遍。dp[1]200是指医院设在1号点全体居民到1号的总距离1号300人×0 3号100人×2 200。对。换根到2号公式dp[2] dp[1] tot - 2*sz[2] 200 400 - 2×100 400。但实际上医院设在2号总距离是300×1 100×1 400。公式算出来是对的和我手算一致。这里2的sz是1002号自己的0人3号100人所以总人口400减去200得增量200正好等于1号方向300人多走一条边300减去3号方向100人少走一条边100净增200。逻辑对上了。再用该公式往3号推dp[3] dp[2] tot - 2*sz[3] 400 400 - 200 600手算300人走2条边600医院在3号确实总代价600。所以该链最优解是医院设在2号总代价400。这也符合直觉——2号居中两边人口差不多。4. 手把手实现完整C代码与四个容易翻车的细节4.1 建树与读入孩子为0数组要清零P1364的节点编号从1开始0表示空孩子。写代码时优先用全局数组因为全局数组自动清零不用手动memset。如果图省事把数组定义在main里一定记得memset(lc, 0, sizeof(lc))否则上一次测试数据可能污染这次结果。我在现场赛就吃过这个亏——本地数组第二次循环时会带着上一组的残留值。另一个容易漏掉的点是邻接表存双向边。如果只存单向的孩子关系Floyd和BFS都会出错。初步判断自己是不是理解对了就看代码里有没有为每条边push两次。4.2 BFS暴力版的vis重置与距离累加BFS暴力版最经典的bug是多次BFS共用同一个dist数组第二次BFS开始时没把dist全部重置成-1。如果只把起点s的距离置0上一次BFS留下的旧距离会被新BFS读到导致累加结果莫名其妙变很小。建议每次BFS都新建一个vectorint dist(N, -1)成本不高但能杜绝这类问题。另外在BFS的累加逻辑里res w[v] * dist[v]这句一定要在入队时做不要在出队时做。出队时做虽然也能算对但会多算一次起点本身起点dist为0不影响逻辑上更容易绕晕还是把维护工作放在入队时最清晰。4.3 树形DP用int还是long long这题n≤100每个节点的居民数我记得上限是100所以tot ≤ 10000dp最大大概在 10^6级别int完全没问题。但做题养成的习惯是树形DP涉及累加时直接开long long。原因很简单题目改一个数据范围int就可能溢出而long long几乎不会。尤其当你从P1364扩展到n10⁵的变体题时这个习惯能帮你少挂一次测试点。4.4 完整可交代码换根DP下面这份代码是正解C17提交即可。#include bits/stdc.h using namespace std; using ll long long; const int N 105; int n; int w[N], lc[N], rc[N]; ll sz[N], dp[N]; ll ans; void dfs1(int u) { sz[u] w[u]; if (lc[u]) { dfs1(lc[u]); sz[u] sz[lc[u]]; dp[u] dp[lc[u]] sz[lc[u]]; } if (rc[u]) { dfs1(rc[u]); sz[u] sz[rc[u]]; dp[u] dp[rc[u]] sz[rc[u]]; } } void dfs2(int u) { ans min(ans, dp[u]); if (lc[u]) { dp[lc[u]] dp[u] sz[1] - 2 * sz[lc[u]]; dfs2(lc[u]); } if (rc[u]) { dp[rc[u]] dp[u] sz[1] - 2 * sz[rc[u]]; dfs2(rc[u]); } } int main() { cin n; for (int i 1; i n; i) { cin w[i] lc[i] rc[i]; } dfs1(1); ans dp[1]; dfs2(1); cout ans endl; return 0; }代码的骨架就是两次DFS。第一次从1号点往下走把整棵树捋一遍得到每个子树的sz和dp第二次再往下走一边走一边用公式推孩子的答案并记录全局最小值。这里有个值得注意的点dfs2里的sz[1]就是总人口tot因为1号点是第一次DFS的根它的sz等于全员。直接用sz[1]写公式比额外开一个tot变量更少出错。4.5 一个关于“二叉树”的提醒我在文章开头提过这题最大的坑是把“二叉树”理解成“只能往子树走”。但正解代码里第二次DFS恰恰是从上往下推的看起来好像也是“只走下边”——为什么不漏关键在dp[u]的定义它不是“以u为根的子树从u向下所有居民的距离”而是整棵树所有居民到u的总距离。第一次DFS算的是局部第二次DFS用公式把父节点的全局值转化成了子节点的全局值。所以即使代码只向下走也没有丢失父方向的信息因为父方向的代价已经藏在dp[u]里了。如果非要检验自己是否理解可以试试把代码里的dfs2改成从任意叶子往上推或者把根从1换成别的节点跑一遍看看最终答案是否一样。理论上换根的结果是等价的这是树形DP“根无关性”的表现。5. 从P1364到一类题带权重心还能怎么考5.1 题型识别什么时候用换根DP刷题多了之后你会发现以下句型都是换根DP的“通缉令”“在树上选一个点使所有点到该点的距离之和最小”“选一个位置建仓库/医院/学校使运输总费用最小”“各点有权值边有权值求Σ(点权×点到选择点的距离)最小”“一条链上开会所有人走路总时间最短”。一旦出现在树上的这类题第一反应就应该是带权重心。数据范围如果n≤2000暴力枚举O(n²)没问题n≤10⁵甚至更大就必须写二次扫描换根。5.2 变式题怎么扩展公式原题边权都等于1所以公式里的系数是1。如果边权不同比如一条从u到v的边有花费cost转移公式变成dp[v] dp[u] (tot - 2 × sz[v]) × cost推导完全一样只是跨过这条边时子树内每个人减少的是cost子树外每个人增加的也是cost。如果点权变成边权每条边本身的通行费用由经过人数决定又可以把边权拆成“边上所有人加起来”的形式再套公式。这些变式在竞赛里都有对应题目理解了P1364原型基本上可以直接套模板改两行。5.3 我的做题习惯与实操建议最后分享几个我的实操习惯希望能少走点弯路。第一先写暴力再写正解对拍验证。对P1364这种小数据题先交一个Floyd或者BFS暴力版本再用正解和它对拍几百组随机数据。随机数据用链、满二叉树、随机二叉树三种形态各生成一些覆盖形态差异。暴力版和正解结果一致才说明换根DP写对了。第二递归深度问题提前预判。P1364的n100递归深度最多100无所谓。但如果复制这个模板去做n10⁵的题链状树会让递归深度爆栈。到时候要么改成显式栈迭代要么在比赛环境下用#pragma comment(linker, /STACK:102400000,102400000)Windows下或者加大栈空间Linux下。更稳妥的做法是练习时就用迭代版DFS或者用std::vector模拟栈来转移树形DP。第三把这类题归类成“换根DP模板”单独记笔记。树上最长路径、树的直径、带权重心、以及“树上选两个点使距离最大/最小”这类题解法套路都很接近第一次DFS求子树的某个量第二次DFS用父节点的答案推导子节点。模板一旦建立再遇到类似的题就非常省力。回到P1364本身这题的真正价值不在于让你AC一道二叉树题而在于让你建立“从暴力到优化”的完整思维路径。先会暴力BFS再会Floyd最后理解换根公式一套流程走下来带权重心这个知识点才算真正扎进脑子里。做题时不妨多问自己一句如果n开大十倍我的代码还能过吗这个问题问多了很多看似“碰运气”的AC就会变成真正有把握的AC。