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; } };