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;
}
}
};
148 ms 26.4 MB
我的CSDN博客地址 https://michael.blog.csdn.net/
长按或扫码关注我的公众号（Michael阿明），一起加油、一起学习进步！
