leedcode : https://leetcode-cn.com/problems/sum-of-distances-in-tree/

思路 : 一遍 dfs 可以求到以任意一个点为根,所有点所含孩子的个数,cnt数组记录
考虑每一条边的贡献时,即此边左右节点数相乘
第2遍 dfs, 以其他点为根时,例如以2为根,先去掉 0-1 边的贡献,再加上以2为根 0和1 对 2 的贡献
代码示例 :
class Solution {
public:
int cnt[10005];
int ans[10005];
vector<int>ve[10005];
int pt[10005];
void dfs(int x){
//printf("++++%d %d
", x, ve[x].size());
pt[x] = 1; cnt[x] = 1;
int sum = 0;
for(int i = 0; i < ve[x].size(); i++){
if (!pt[ve[x][i]]) {
dfs(ve[x][i]);
sum += cnt[ve[x][i]];
//printf("+++ %d %d %d %d
", x, ve[x][i], cnt[x], cnt[ve[x][i]]);
ans[0] += cnt[x] * cnt[ve[x][i]];
}
}
cnt[x] += sum;
}
void dfs2(int x, int N){
pt[x] = 1;
for(int i = 0; i < ve[x].size(); i++){
if (!pt[ve[x][i]]) {
ans[ve[x][i]] = (ans[x]-cnt[ve[x][i]])+(N-cnt[ve[x][i]]);
dfs2(ve[x][i], N);
}
}
}
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
vector<int> sum;
for(int i = 0; i < N-1; i++){
int u = edges[i][0];
int v = edges[i][1];
ve[u].push_back(v);
ve[v].push_back(u);
}
memset(cnt, 0, sizeof(cnt));
memset(pt, 0, sizeof(pt));
dfs(0);
memset(pt, 0, sizeof(pt));
dfs2(0, N);
for(int i = 0; i < N; i++){
sum.push_back(ans[i]);
}
return sum;
}
};