CF1324F Maximum White Subtree

题目

CF1324F Maximum White Subtree

分析

换根 \(dp\) 。

代码

#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define ull unsigned long long
#define inc(x,y,mod) (((x)+(y))>=(mod)?(x)+(y)-(mod):(x)+(y))
#define dec(x,y,mod) ((x)-(y)<0?(x)-(y)+(mod):(x)-(y))
#define rep(i,x,y) for(int i=(x);i<=(y);i++)
#define dep(i,y,x) for(int i=(y);i>=(x);i--)
#define pb push_back
const int N=2e5+5,M=2e5+5,MOD=1e9+7;
int n,dp[N],Ans[N];
bool a[N];
vector<int>G[N];
void dfs(int x,int fa){ 
	dp[x]=a[x]?1:-1;
	for(int i=0;i<G[x].size();i++){
		int y=G[x][i];
		if(y==fa)continue;
		dfs(y,x);
		dp[x]+=max(0,dp[y]);
	}
	return ;
}
void dfs1(int x,int fa){
	Ans[x]=dp[x];
	for(int i=0;i<G[x].size();i++){
		int y=G[x][i];
		if(y==fa)continue;
		dp[x]-=max(0,dp[y]);dp[y]+=max(0,dp[x]);
		dfs1(y,x);
		dp[y]-=max(0,dp[x]);dp[x]+=max(0,dp[y]);
	}
}
signed main(){
	ios::sync_with_stdio(false);
	cin>>n;
	rep(i,1,n) cin>>a[i];
	rep(i,1,n-1){int x,y;cin>>x>>y;G[x].pb(y);G[y].pb(x);}
	dfs(1,0);
	dfs1(1,0);
	rep(i,1,n) cout<<Ans[i]<<" ";
	return 0;
}