codeforces round 627 div3 [树上dp]F. Maximum White Subtree

#include<iostream>
#include<vector>
#define pb push_back
 
using namespace std;
const int N = 2e5 + 5;
vector<int> g[N];
int dp[N],a[N];
const int INF = 0x3f3f3f3f;
int dfs(int u,int pre){
	int ans = (a[u] == 1)?1 : -1;	
	for(unsigned int i = 0;i < g[u].size();i++){
		int v = g[u][i];
		if(v != pre)ans += dfs(v,u);
	}
	dp[u] = ans;
	return max(ans,0);
}
void dfs1(int u,int pre){
	for(unsigned int i = 0;i < g[u].size();i++){
		int v = g[u][i];
		if(v != pre && dp[u] - max(dp[v],0) > 0){
			dp[v] = dp[u] - max(dp[v],0) + dp[v];
		}
		if(v != pre)dfs1(v,u);
	}
}

int main(){
	int n;
	cin>>n;
	for(int i = 1;i <= n;i++)cin>>a[i],dp[i] = -INF;
	for(int i = 1;i <= n - 1;i++){
		int u,v;
		cin>>u>>v;
		g[u].pb(v);
		g[v].pb(u);
	}
	dfs(1,-1);
	dfs1(1,-1); 
	for(int i = 1;i <= n;i++)
		cout<<dp[i]<<" ";
	return 0;
}