Maximum White Subtree CodeForces

#include<queue>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm> 
using namespace std;
int read()
{
    int res=0,ch,flag=0;
    if((ch=getchar())=='-')             //判断正负
        flag=1;
    else if(ch>='0'&&ch<='9')           //得到完整的数
        res=ch-'0';
    while((ch=getchar())>='0'&&ch<='9')
        res=res*10+ch-'0';
    return flag?-res:res;
}
const int N=5e5+100;
int n;
int a[N];
int idx,h[N],e[N],ne[N];
void add(int a,int b)
{
	e[idx]=b;
	ne[idx]=h[a];
	h[a]=idx++;
}
//f[x]:在以x为根的子树内，所有包含点x的联通子图中，白点数减去黑点数的最大值
//1为所有的根，从1开始往下搜
//不能往回遍历其父节点 
int f[N];
//f[1]==g[1]
//然后往下遍历，画图理解 
int g[N];
void dfs(int u,int fa)
{
	if(a[u]==1)
		f[u]++;
	else
		f[u]--;
	for(int i=h[u];i!=-1;i=ne[i])
	{
		int v=e[i];
		if(v==fa)
			continue;
		dfs(v,u);
		if(f[v]>0)
			f[u]+=f[v];
	}
}
bool vis[N];
void bfs()
{
	queue<int>q;
	q.push(1);
	g[1]=f[1];
	vis[1]=1;
	while(q.size())
	{
		int u=q.front();
		q.pop();
		for(int i=h[u];i!=-1;i=ne[i])
		{
			int v=e[i];
			if(vis[v])
				continue;
			g[v]=f[v]+max(0,g[u]-max(0,f[v]));
			vis[v]=1;
			q.push(v);
		}
	}
}
int main()
{
	memset(h,-1,sizeof h);
	n=read();
	for(int i=1; i<=n; i++)
		a[i]=read();
	for(int i=1; i<n; i++) {
		int u=read(),v=read();
		add(u,v),add(v,u);
	}
	dfs(1,-1);
	bfs();
	for(int i=1; i<=n; i++)
		printf("%d ",g[i]);
	cout<<endl;
	return 0;
}