困难：834. 树中距离之和
解题思路
此题在我收藏夹躺了好久了，今天终于轮到练习它了。不久前在第 328 场周赛中，也遇到过一道类似的困难题。
据说这两题要用树状 dp来解，一直觉得很高大上。其实最开始在对这两道题冥思苦想之际，我就隐约意识到可以用动态规划来解，但是实在想不到该怎么做状态转移。今天详细学习了官解，发现其实挺简单的。
树状 dp的特点，就是动态规划的每个子问题，一一对应于一颗子树。可以从由单个叶子节点构成的子树开始，自底向上进行规划；也可以从根节点开始，自上而下地、使用带缓存的方式进行规划。后者本质就是一个对整棵树 dfs 的过程。
树状 dp，时间复杂度为 。
有一类与无根树相关的题目，根节点可以任意选择，需要对每一种可能的根节点，都运行一次树状 dp，故整体时间复杂度为。可以通过反复换根的方式，将这种题目的时间复杂度度优化到 。这个换根的过程，也采取 dfs 的方式。
本题就属于这种类型。需要先随意选择一个根，运行一次树状 dp，然后额外运行一次 bfs 进行换根。也就是，总共进行两次 dfs。
我看到很多扣友将这种两次 dfs 的思路，叫做换根 dp，很形象，但是我暂时没有找到这个术语的权威出处，叫着有点心虚。如果哪位小伙伴知道出处，可以告我一声🤝。
官解挺容易理解的，只要看懂下图就没啥大问题：
红框里的状态转移方程有点小问题，一开始看的我有点头晕，应该是：
而不是：
编码细节
理解了思路，让后我开始尝试写代码，发现还是有很多细节的。
树状 dp，我感觉还是采用递归的方式比较好写，如果用自底向上的方式，需要先进行层序遍历（bfs），对节点按层分组，然后从最底层开始进行动态规划。
为了循序渐进，我先写了一个进行 n 次树状 dp 的暴力版本，这样做果然超时了，但是我的目的是为了分步进行练习，找找写树状 dp 的感觉。毕竟一口真的很难吃成个胖子。
确保暴力版本正确后，我开始尝试换根，这个过程调试起来很痛苦，折腾了很久。一定要时刻记住换根的本质，防止神志混乱：子节点的子节点多了一个，父节点的子节点少了一个。
换根的过程，虽然也是 dfs，但和动态规划毫无关系，不要写混了。
换根，貌似必须按照 dfs 的方式来，我一开始尝试用 bfs，结果发现很难写。用 bfs 每次换完根后，可以通过回溯，将根换回来，方便兄弟节点继续和父节点交换。也许 bfs 也可以，但是我没写出来。
代码
暴力版本：n 次树状 dp
注意，暴力版会超时，仅做练习之用。
class Solution {
public int[] sumOfDistancesInTree(int n, int[][] edges) {
List<Integer>[] adt = new List[n];
for(int i = 0; i < n; i++)adt[i] = new ArrayList<Integer>();
for(int[] e: edges){
adt[e[0]].add(e[1]);
adt[e[1]].add(e[0]);
}
int[] ret = new int[n];
for(int i = 0; i < n; i++){
//System.out.printf("%s--------------------\n", i);
// 以 i 为树的根节点，进行自上而下带缓存的的动态规划
int[] dp = new int[n]; // dp[i] 代表：对于以节点 i 为根的子树，i 到所有后代的距离的和
int[] size = new int[n]; // size[i] 代表以节点i 为根的子树中，所有节点的数量
helper(n, adt, i, i, dp, size);
ret[i] = dp[i];
}
return ret;
}
/**
用动态规划，递归地求出以 u 为根节的子树中，u 到所有后代节点的距离之和。pa 代表 u 的直接父节点
调用一次 helper，会求出 dp[u] 和 size[u]
*/
private void helper(int n, List<Integer>[] adt, int pa, int u, int[] dp, int[] size){
if(size[u] > 0)return;
size[u]++;
for(int v: adt[u]){
if(v == pa)continue;
helper(n, adt, u, v, dp, size);
dp[u] += dp[v] + size[v];
size[u] += size[v];
}
//System.out.printf("%s: %s, %s\n", u, size[u], dp[u]);
}
}树状 dp + dfs 换根
换根的逻辑，被单独提取成了 swap 方法，可以让代码简洁一点：
class Solution {
public int[] sumOfDistancesInTree(int n, int[][] edges) {
List<Integer>[] adt = new List[n];
for(int i = 0; i < n; i++)adt[i] = new ArrayList<Integer>();
for(int[] e: edges){
adt[e[0]].add(e[1]);
adt[e[1]].add(e[0]);
}
int[] ret = new int[n];
int[] dp = new int[n]; // dp[i] 代表以 i 为根的子树中， i 到所有后代节点的距离之和。i 的父节点可以有多种情况
int[] size = new int[n]; // size[i] 代表以节点 i 为根的【子树】中，所有节点的数量
helper(n, adt, -1, 0, dp, size);
ret[0] = dp[0];
dfs(adt, -1, 0, ret, dp, size);
return ret;
}
/**
用动态规划，递归地求出以 u 为根节的子树中，u 到所有后代节点的距离之和。pa 代表 u 的直接父节点
调用一次 helper，会求出 dp[u] 和 size[u]
*/
private void helper(int n, List<Integer>[] adt, int pa, int u, int[] dp, int[] size){
if(size[u] > 0)return;
size[u]++;
for(int v: adt[u]){
if(v == pa)continue;
helper(n, adt, u, v, dp, size);
dp[u] += dp[v] + size[v];
size[u] += size[v];
}
//System.out.printf("%s: %s, %s\n", u, size[u], dp[u]);
}
// 用回溯法不断换根
private void dfs(List<Integer>[] adt,int pa, int u, int[] ret, int[] dp, int[] size){
for(int v: adt[u]){
if(v == pa)continue;
swap(u, v, dp, size);
ret[v] = dp[v];
//System.out.printf("%s %s, %s\n", u, v, ret[v]);
dfs(adt, u, v, ret, dp, size);
swap(v, u, dp, size);
}
}
private void swap(int u, int v, int[] dp, int[] size){
dp[u] -= dp[v] + size[v];
size[u] -= size[v];
dp[v] += dp[u] + size[u];
size[v] += size[u];
}
}举一反三：更多树状 dp 题目（含代码）
困难：2538. 最大价值和与最小价值和的差值
如果这道题比较痛苦，可以跳过，先做后面的中等题 1245。
这道题有两点需要注意：
- 确保每次换根在 内完成，需要点技巧：对要被替换的某子树根节点
u，先遍历其邻接点一遍，求出 dp 值的最大之和次大值。如果不注意，会喜提超时，详见下面的错误版本①和正确版本②。 - 可以换一种思路，不用换根，直接进行一次 树状dp 即可。这个思路来自灵茶山艾府大佬的视频题解。我一开始怎么也理解不了灵佬的思路，后来发现灵佬没有换根，是一种不同的思路。
下面的四个版本中，后三个都能良好的工作，版本③和版本④思路是一样的。可以看到，虽然都用了树形 dp，但是定义的子问题有差别，所以状态转移的过程，也有差异。
/*
单次换根没有在 O(1)内完成，导致超时。
834 题可以将换根提取到一个简单的 swap 方法中，比较简洁，但是这道题不可以。
通过额外遍历一次邻接点，求出最大值和次大值，确保了O(1)的均摊单次换根代价。
*/
/*
单次换根没有在 O(1)内完成，导致超时。
834 题可以将换根提取到一个简单的 swap 方法中，比较简洁，但是这道题不可以。
通过额外遍历一次邻接点，求出最大值和次大值，确保了O(1)的均摊单次换根代价。
*/
class Solution {
private List<Integer>[] adt; // 邻接表
private long max = 0;
private long[] dp; // dp[i] 代表以 i 为根节点的子树中，i 最大的【开销】
private int[] price;
public long maxOutput(int n, int[][] edges, int[] price) {
this.adt = new List[n];
this.max = 0;
this.dp = new long[n];
this.price = price;
// 构造邻接表
for(int i = 0; i < n; i++)adt[i] = new ArrayList<>();
for(int[] e: edges){
adt[e[0]].add(e[1]);
adt[e[1]].add(e[0]);
}
// 以节点 0 为根，运行一次自上而下的带缓存的树形动态规划
helper(-1, 0);
max = dp[0];
// 用回溯法不断换根，求出 max
dfs(-1, 0);
return max;
}
// 自上而下的带缓存的树形动态规划
// pa代表 u 的父节点
private void helper(int pa, int u){
for(int v: adt[u]){
if(v == pa)continue;
helper(u, v);
dp[u] = Math.max(dp[u], dp[v] + price[v]);
}
System.out.printf("%s,%s\n", u, dp[u]);
}
private void dfs(int pa, int u){
for(int v: adt[u]){
if(v == pa)continue;
swap(u, v);
max = Math.max(dp[v], max);
System.out.printf("换根：%s,%s, %s, %s, %s\n", u, v, dp[u], dp[v], max);
dfs(u, v);
swap(v, u);
}
}
private void swap(int u, int v){
dp[u] = 0;
for(int x: adt[u]){
if(x == v)continue;
dp[u] = Math.max(dp[u], dp[x] + price[x]);
}
dp[v] = Math.max(dp[v], dp[u] + price[u]);
}
}中等：1245. 树的直径
不知道为啥，这道题虽然曾出现在周赛中，但是现在已经是一道会员题了，截个图：
这道题可以采用三种方式求解
- 树形 dp + 换根，换根时，对任一颗子树的根节点 ，需要遍历其邻接点，并从中求出 dp 值最大和次大的子节点。这个思路类似于前面2538. 最大价值和与最小价值和的差值中的代码版本②。如果不这样做，在换根时遍历一遍邻接点，换根会超时。
- 一次树形 dp 即可。代价是，做动态规划时，要同时维护每颗子树的根节点 能达到的最大路径和次大路径。所有子树中，最大路径和次大路径的和的最大值，就是答案。
- 改进一下思路 2，详见下面的代码版本③。
下面的三个版本，都能正常工作。但是状态转移的方式，各不相同。
class Solution {
public int treeDiameter(int[][] edges) {
int n = edges.length + 1;
int[] dp = new int[n]; // dp[i] 代表以 i 为根节点的子树中，i 到其后代节点的最长路径
// 生成邻接表
List<Integer>[] adt = new List[n];
for(int i = 0; i < n; i++) adt[i] = new ArrayList<Integer>();
for(int[] e: edges){
adt[e[0]].add(e[1]);
adt[e[1]].add(e[0]);
}
dfs(adt, -1, 0, dp);
int[] max = new int[]{dp[0]};
dfs2(adt, -1, 0, dp, max);
return max[0];
}
private void dfs(List<Integer>[] adt, int pa, int u, int[] dp){
for(int v: adt[u]){
if(pa == v)continue;
dfs(adt, u, v, dp);
dp[u] = Math.max(dp[u], dp[v] + 1);
}
// System.out.printf("%s, %s\n", u, dp[u]);
}
private void dfs2(List<Integer>[] adt, int pa, int u, int[] dp, int[] max){
int m1 = 0, m2 = 0; //最大值和次大值
for(int v: adt[u]){
int val = dp[v] + 1;
if(m1 <= val){
m2 = m1;
m1 = val;
}else{
m2 = Math.max(m2, val);
}
}
for(int v: adt[u]){
if(pa == v)continue;
int baku = dp[u], bakv = dp[v];
if(dp[v] + 1 == dp[u])dp[u] = m2;
dp[v] = Math.max(dp[v], dp[u] + 1);
max[0] = Math.max(max[0], dp[v]);
dfs2(adt, u, v, dp, max);
dp[u] = baku;
dp[v] = bakv;
}
}
}中等：310. 最小高度树
做过了前面三道题，这道题，用树状 dp + 换根，可以说非常轻松，思路也很自然。这里在换根时，依然采用了求最大和次大值的技巧。
这道题还有很多别的解法，比如官解给出了三种解法，前两种比较数学化，看的头疼，暂时放弃了。官解的第三种比较有趣，但是官解为其打了个“拓扑排序”的标签，我觉得这个标签不合适，虽然过程有点类似用 bfs 求拓扑排序，但是二者其实风马牛不相及。官解三非常直观。
宫水三叶的题解【宫水三叶】树形 DP 运用题中，使用了一种没有换根，但是对我来说略难懂的思路，可作为参考。
class Solution {
private List<Integer>[] adt; // 邻接表
private int[] dp; // dp[i] 代表以 i 为根的子树的高度。注意，整棵树的根，是动态可变的。
private int[] dp2; // dp2 [i] 代表整棵树以 i 根时，树的高度
private int min;
public List<Integer> findMinHeightTrees(int n, int[][] edges) {
dp = new int[n];
dp2 = new int[n];
min = Integer.MAX_VALUE;
adt = new List[n];
Arrays.setAll(adt, value -> new ArrayList<>());
for (int[] e : edges) {
adt[e[0]].add(e[1]);
adt[e[1]].add(e[0]);
}
// 进行一次树状 dp，以 0 为根节点
helper(-1, 0);
dp2[0] = dp[0];
min = Math.min(dp[0], min);
// 换根
dfs(-1, 0);
List<Integer> ret = new ArrayList<>();
for(int i = 0; i < n; i++){
if(dp2[i] == min)ret.add(i);
}
return ret;
}
/**
* 换根
* @param pa u 的父节点
* @param u
*/
private void dfs(int pa, int u) {
int m1 = 0, m2 = 0; // 最大和次大
for (int v : adt[u]) {
int val = dp[v] + 1;
if(m1 <= val){
m2 = m1;
m1 = val;
}else{
m2 = Math.max(m2, val);
}
}
for (int v : adt[u]) {
if(v == pa)continue;
int baku = dp[u], bakv = dp[v];
if(dp[v] + 1 == dp[u])dp[u] = m2;
dp[v] = Math.max(dp[u] + 1, dp[v]); // 求出以 v 为整颗树的根时，树的高度
min = Math.min(min, dp[v]);
dp2[v] = dp[v];
// System.out.printf("换根%s, %s,%s, %s, %s\n",u, v, dp[v], min, ret);
dfs(u, v);
dp[u] = baku;
dp[v] = bakv;
}
}
/*
* 自上而下带缓存的动态规划，pa 是 u 的父节点
* */
private void helper(int pa, int u) {
for (int v : adt[u]) {
if(pa == v)continue;
helper(u,v);
dp[u] = Math.max(dp[u], dp[v] + 1);
}
//System.out.printf("%s, %s\n", u, dp[u]);
}
}总结
解完这 4 道题，对树状 dp 感觉没那么神秘和恐怖了。其实不管有没有“树状”这个修饰词，dp 最让人害怕的，还是怎么想出状态转移方程。不同的状态转移思路，对应不同的解法。2538. 最大价值和与最小价值和的差值这道题，灵茶山艾府大佬的转移方式，我现在还感觉恍恍惚惚、叹为观止🐒
这四道题有个共同点，那就是都可以转换为求“最长路径”。我的刷题量有限，可能并不是所有的树状 dp 问题都有这个共性，这个只能以后进一步探索了。
换根这个技巧，掌握了以后比较容易上手；追求不换根，反而比较难。至少对这四道题是如此。
搞完这几道题，我的图论之旅总算快谢幕了。接下来，该系统地深扒一下变幻莫测动态规划了。😵😵😵
彩蛋 🌈🌈🌈
今天颇费了一番功夫，然后解锁了一个小技能。写题解时，如果有多段不同的代码，可以让它们以 tab 标签的形式展示，并且可以自定义 tab 文字，从而让题解更紧凑：
这并不是 markdown 的原生语法，只适用于 leecode。
有没有很神奇？点个赞👍，我来告诉你怎么解锁这个小技能，嘻嘻~
