给定一个无向、连通的树。树中有 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
著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。
ans[i] : i节点到其他所有节点的距离
sum[i] : 以i节点为root的子树里的所有节点到i的距离
cnt[i] : 以i节点为root的子树的大小,包括i自己
将问题拆分:对于两个相邻节点A和B,将树拆分为两个子树,根节点分别为A和B,ans(A) = sum(A) + sum(B) + cnt(B);
*如何理解cnt(b)?因为A、B相邻,已经求出B到其子树里所有节点的距离,对A来说,对每个节点都+1,就变成了A到B子树里所有节点的距离。
同理,ans(B) = sum(B) + sum(A) + cnt(A);
由此我们得到:
ans(A) = sum(A) + sum(B) + cnt(B);
ans(B) = sum(B) + sum(A) + cnt(A);
ans(A) = ans(B) - cnt(A) + cnt(B) = ans(B) - cnt(A) + (N - cnt(A));
因此,对于根节点root的任意子节点child,ans(child) = ans(root) - cnt(child) + N - cnt(child);
得到root的答案就可以DFS递归得到其他所有节点的答案。(这里需要一个DFS)
class Solution {
public:
vector <int> ans;
vector <int> cnt;
vector <vector <int>> G;
int n;
void dfs(int now, int fa) {
for (int i = 0; i < G[now].size(); i++) {
int v = G[now][i];
if (v != fa) {
dfs(v, now);
cnt[now] += cnt[v];
ans[now] += ans[v] + cnt[v];
}
}
}
void dfs2(int now, int fa) {
for (int i = 0; i < G[now].size(); i++) {
int v = G[now][i];
if (v != fa) {
ans[v] = ans[now] - cnt[v] + (n - cnt[v]);
dfs2(v, now);
}
}
}
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
if (N == 0 || edges.empty()) {
return {0};
}
n = N;
G.resize(N);
ans.resize(N);
cnt.resize(N);
for (int i = 0; i < n; i++) {
cnt[i] = 1;
}
for (auto edge : edges) {
int u = edge[0];
int v = edge[1];
G[u].push_back(v);
G[v].push_back(u);
}
//第一次dfs求得节点0的正确答案
//第一次dfs后的ans[i]表示节点i到其所有子节点的距离之和
dfs(0, -1);
//第二次dfs通过节点0的答案求得其他节点的正确答案
//第二次dfs后的ans[i]表示节点i到其所有 节点的距离之和
dfs2(0 , -1);
return ans;
}
};