Atcoder ABC 220F（换根dp）

https://atcoder.jp/contests/abc220/tasks/abc220_f

换根dp。

首先从根节点dfs一遍，统计出所有子树的大小。

然后求出根节点的答案。

最后换根dp，推出式子：f[v]=f[u]+n-sz[v]*2;
注意是整棵树的大小减去子树大小

#include<cstdio>
#include<iostream>
#include<cstring>
using namespace std;
const int N=200005;
typedef long long ll;
int n;
int tot;
ll ans;
int head[N];
ll sz[N];
ll f[N];
struct node{
	int to,next;
}edge[N<<1];
void add(int u,int v){
	edge[tot].to=v;
	edge[tot].next=head[u];
	head[u]=tot++;
}
void dfs(int u,int fa){
	for(int i=head[u];i!=-1;i=edge[i].next){
		int v=edge[i].to;
		if(v==fa) continue;
		dfs(v,u);
		sz[u]+=sz[v];
	}
}
void dfs2(int u,int fa,int step){
	ans+=step;
	for(int i=head[u];i!=-1;i=edge[i].next){
		int v=edge[i].to;
		if(v==fa) continue;
		dfs2(v,u,step+1);
	}
}
void dp(int u,int fa){
	for(int i=head[u];i!=-1;i=edge[i].next){
		int v=edge[i].to;
		if(v==fa) continue;
		f[v]=f[u]+n-sz[v]*2;//整棵树的大小减去子树大小 
		dp(v,u);
	}
}
int main(){
	memset(head,-1,sizeof(head));
	scanf("%d",&n);
	for(int i=1;i<=n;i++) sz[i]=1;
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%d%d",&u,&v);
		add(u,v); add(v,u);
	}
	dfs(1,0);
	dfs2(1,0,0);
	f[1]=ans;
	dp(1,0);
	for(int i=1;i<=n;i++) printf("%lld
",f[i]);
	return 0;
}