LeetCode 834. 树中距离之和（树上DP）*

文章目录

- 1. 题目
- 2. 解题

1. 题目

给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边。

第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

示例 1:
输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 
如下为给定的树的示意图：0/ \
1   2/|\3 4 5我们可以计算出 
dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 
也就是 1 + 1 + 2 + 2 + 2 = 8。 
因此，answer[0] = 8，以此类推。
说明: 1 <= N <= 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

2. 解题

抄的官方答案

class Solution {vector<int> ans, son, dp;vector<vector<int>> g;
public:vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {ans.resize(N);son.resize(N);dp.resize(N);g.resize(N);for(auto& e : edges){g[e[0]].push_back(e[1]);g[e[1]].push_back(e[0]);}dfs1(0, -1);//dp 初始化dfs2(0, -1);//dp 递推return ans;}void dfs1(int u, int f){	//求子树包含的节点数量// 第一次dfs 得到的每个节点子树节点到该节点的距离和son[u] = 1;dp[u] = 0;for(auto v : g[u]){if(v == f)continue;dfs1(v, u);dp[u] += dp[v]+son[v];//子节点的距离和 + 边被占用次数son[u] += son[v];}}void dfs2(int u, int f){ans[u] = dp[u];for(auto v : g[u]){if(v == f)continue;int dpu = dp[u], dpv = dp[v];//记录状态int sonu = son[u], sonv = son[v];dp[u] -= dp[v]+son[v];//跟换根节点为 v ，u 变成子节点son[u] -= son[v];dp[v] += dp[u]+son[u];// v 加上 u 的数据son[v] += son[u];dfs2(v, u);dp[u] = dpu, dp[v] = dpv;//回溯状态son[u] = sonu, son[v] = sonv;}}
};