题目描述
给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
样例
输入: 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,以此类推。
树形dp $O(n)$
维护一个子树u中节点个数cnt[u]及所有节点到u的距离和f[u],先dfs()将f[0]算出来,再逐个将其他节点作为根节点计算f[u]作为res[u]
f[u] = 子树v中子节点到子树根节点v的距离he + 1×子树的节点个数(子树中每个节点到u的距离多了v→u这条边)
$f[u] = \sum([v]+cnt[v])$
换根后
先将f[u]和cnt[u] 减去 v的贡献值
$f[u] -= f[v]+cnt[v];
cnt[u]-=cnt[v];$
再将f[u]和cnt[v] 加上 此时u作为v子节点的值
$f[v] += f[u]+cnt[u];
cnt[v]+=cnt[u];$
C++ 代码
const int M=20010;
class Solution {
public:
int h[M],e[M],ne[M],idx=0;
int cnt[M],f[M];
vector<int> res;
void add(int a,int b)
{
e[idx] = b;
ne[idx] = h[a];
h[a] = idx++;
}
void dfs2(int u,int fa)
{
res[u] = f[u];
for(int i=h[u];i!=-1;i=ne[i])
{
int v = e[i];
if(v==fa) continue;
int pu = f[u], pv = f[v];
int su = cnt[u], sv = cnt[v];
f[u] -= f[v]+cnt[v];
cnt[u]-=cnt[v];
f[v] += f[u]+cnt[u];
cnt[v]+=cnt[u];
dfs2(v,u);
// 把v的更新完后 继续对u的另一个子节点 v‘更新时要还原u到之前的状态
f[u] = pu,f[v]=pv;
cnt[u]=su,cnt[v]=sv;
}
}
void dfs(int u,int fa)
{
cnt[u]=1;
f[u]=0;
for(int i=h[u];i!=-1;i=ne[i])
{
int v=e[i];
if(v==fa)continue;
dfs(v,u);//写错成u,v了
f[u] += f[v]+cnt[v];
cnt[u] += cnt[v];
}
}
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
memset(h,-1,sizeof h);
for(auto t:edges)
{
add(t[0],t[1]);
add(t[1],t[0]);
}
res.resize(N,0);
dfs(0,-1);
dfs2(0,-1);
return res;
}
};