ICPC NEAU Programming Contest 2020 G. 选根 (换根DP)

解法： 换根dp

Code:

const int MX = 2e5 + 7;


int v[MX];
ll sum[MX],dp[MX],ans[MX];
int dep[MX];
vector<int>e[MX];
void init(int n){
	for(int i = 1;i <= n;++i){
		sum[i] = dp[i] = dep[i] = 0;
		e[i].clear();
	}
}
void pre_dfs(int u,int fa){
	dep[u] = dep[fa] + 1;
	sum[u] = 1ll * v[u];
	dp[u] = 1ll * v[u] * dep[u];
	for(auto v:e[u]){
		if(v == fa) continue;
		pre_dfs(v,u);
		sum[u] += sum[v];
		dp[u] += dp[v];
	}
}
void dfs(int u,int fa){
	ans[u] = dp[u];
	for(auto v: e[u]){
		if(v == fa) continue;
		ll tmpsumu = sum[u], tmpsumv = sum[v], tmpv = dp[v], tmpu = dp[u];
		sum[u] -= sum[v];sum[v] += sum[u];  
		dp[u] -= tmpv;dp[u] += sum[u];
		dp[v] -= tmpsumv;dp[v] += dp[u];
		dfs(v,u);
		sum[u] = tmpsumu, sum[v] = tmpsumv;dp[u] = tmpu, dp[v] = tmpv;
	}
}
void solve(int n){
	pre_dfs(1,0);
	dfs(1,0);
	for(int i = 1;i <= n;++i){
		printf("%lld ", ans[i]);
	}
	printf("\n");
}
int main(){
	int T;scanf("%d",&T);
	while(T--){
		int n;scanf("%d",&n);init(n);
		for(int i = 1;i <= n;++i) scanf("%d",&v[i]);
		for(int i = 1;i < n;++i){
			int u,v;scanf("%d %d",&u,&v);
			e[u].pb(v);e[v].pb(u);
		}
		solve(n);
	}
	return 0;
}