E. Black and White Tree

Black and White Tree
题意：n个节点的树，至少有两个黑色的节点。你可以选择任意一个黑色节点向它走一步。对于任意两个相邻的操作不能选择同一个黑色的节点。 问从某个结点出发能否到达黑色结点
思路：黑色结点 和黑色结点相邻的一定可以，或者 某个结点的子节点可以且子树中有2个以上黑色结点

//#pragma GCC optimize(2)
//#pragma GCC optimize(3,"Ofast","inline")
#include<bits/stdc++.h>
#define int long long
#define fi first
#define se second
#define pb push_back
#define pii pair<int,int>
#define IOS ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
using namespace std;
const int inf=8e18;
const int maxn=3e5+100;
vector<int>g[maxn];
int ans[maxn];
int a[maxn],n,son[maxn],sum;
void dfs(int x,int f)
{
	if(a[x])son[x]=1;
	for(auto it:g[x])
	{
		if(it==f)continue;
		dfs(it,x);
		son[x]+=son[it];
		if(son[it]>=2&&ans[it])ans[x]=1;
	}
}
void dfs2(int x,int f)
{
	if(sum-son[x]>=2&&ans[f])ans[x]=1;
	for(auto it:g[x])
	{
		if(it==f)continue;
		dfs2(it,x);
	}
}
signed main()
{
	cin>>n;
	for(int i=1;i<=n;i++)
	{
		cin>>a[i];
	}
	for(int i=1;i<n;i++)
	{
		int x,y;
		cin>>x>>y;
		g[x].pb(y);
		g[y].pb(x);
	}
	for(int i=1;i<=n;i++)
	{
		if(a[i])
		{
			ans[i]=1;
			sum++;
			for(auto it:g[i]) ans[it]=1;
		}
	}
	dfs(1,0);
	dfs2(1,0);
	for(int i=1;i<=n;i++)
	{
		cout<<ans[i]<<" ";
	}
}