Solution #

The solution consists of two DFS, first DFS is to calculate $sub_i$: the max difference of the subgraph in subtree of $i$. Note that we treat the graph as a rooted tree, so the subtree means the part that is lower than node $i$. This is pretty naive DFS.

The second dfs is to find the answer of each vertex. Since we only considered the contribution of the subtree, we need to add the contribution of other part the graph. How to find this contribution? We always get the answer of higher nodes first. Let $v$ be the children of node $i$, if $sub_v>0$, the contribution is $ans_i-sub_v$ since we don’t want to count the contribution of subtree twice. If $sub_v\leq 0$, we don’t need to subtract $sub_v$ since we didn’t count it in $ans_i$. If $ans_i-sub_v<0$, it’s useless for $ans_v$. Thus, the contribution of other part is $\max(ans_i-\max(sub_v,0),0)$.

Code #

#include <bits/stdc++.h>
 
#define forn(i, n) for (int i = 0; i < int(n); ++i)
#define for1(i, n) for (int i = 1; i <= int(n); ++i)
 
using namespace std;
typedef long long ll;
 
const int N=2e5+5;
vector<int> G[N];
int ans[N],a[N],dp[N];
void dfs1(int u,int fa){
	dp[u]=a[u];
	for(auto it:G[u]){
		if(it!=fa){
			dfs1(it,u);
			dp[u]+=max(0,dp[it]);
		}
	}
}
void dfs2(int u,int fa,int pd){
	ans[u]=dp[u]+pd;
	for(auto v:G[u]){
		if(v!=fa){
			dfs2(v,u,max(ans[u]-max(dp[v],0),0));
		}
	}
}
int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	int n;
	cin>>n;
	for1(i,n) {
		cin>>a[i];
		if(!a[i]) a[i]=-1;
	}
	forn(i,n-1){
		int x,y;
		cin>>x>>y;
		G[x].pb(y);
		G[y].pb(x);
	}
	dfs1(1,-1);
	dfs2(1,-1,0);
	for1(i,n) cout<<ans[i]<<' ';
	return 0;
}