简单换根DP
其实就是看好变化量,然后让父亲更新儿子就好了~
上图2当根节点的时候,ans[2] = ans[0] -sz[2]+n-sz[2];
class Solution {
public:
vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
vector<vector<int>>g(n);
for(auto &edge:edges){
int a = edge[0],b = edge[1];
g[a].push_back(b),g[b].push_back(a);
}
using ll = long long;
vector<int>ans(n);
vector<ll>sz(n+10);
function<void(int,int,int)>dfs=[&](int u,int father,int dist){
sz[u] = 1;
ans[0]+=dist;
for(auto &v:g[u]){
if(v==father)continue;
dfs(v,u,dist+1);
sz[u]+=sz[v];
}
};
function<void(int,int)>down=[&](int u,int father){
for(auto &v:g[u]){
if(v==father)continue;
ans[v] = ans[u]-sz[v]+n-sz[v];
down(v,u);
}
};
dfs(0,-1,0);
down(0,-1);
return ans;
}
};