1626E - Black and White Tree - CodeForces Solution

#include <bits/stdc++.h> 
using namespace std;
#define ll long long
#define ld long double
#define endl "\n"
typedef vector<ll> vll;
#define pb              push_back 
#define fori(n)         for (ll i=0; i<n; i++) 
#define forj(n)         for (ll j=0; j<n; j++) 
#define fork(n)         for (ll k=0; k<n; k++) 
#define forn(n)         for (ll i=1; i<=n; i++) 
#define forn2(n)        for (ll j=1; j<=n; j++) 
#define forn3(n)        for (ll k=1; k<=n; k++) 
#define Sort(a)         sort(a.begin(),a.end())
const int N=3e5+5;
vll a(N),ans(N),ct(N),black(N),dp(N);
vll g[N];

void dfs(ll v, ll p){
    
    for(auto c:g[v]){
        
        if(a[c]==1){ans[v]=1;}
        if(c==p)continue;
        if(a[c]==1){ans[v]=1;black[v]++;}
        dfs(c,v);
        dp[v]+=dp[c];
        
        if(dp[c]>=2&&black[c])ct[v]++;
        ct[v]+=ct[c];
        if(a[v]==1){ans[c]=1;}
        
    }
}
void dfs2(ll v, ll p){
    if(ct[v])ans[v]=1;
    for(auto c:g[v]){
        
        if(a[c]==1)ans[v]=1;
        if(c==p)continue;

        ll x1=ct[v],x3=ct[c];
        ct[v]-=ct[c];
        ct[c]=x1;
        ll f=0;
        
        ll x=dp[v],x2=dp[c];
        dp[v]-=dp[c];
        dp[c]=x;
        if(a[p]==1)black[v]++;
        if(dp[v]>=2&&black[v]){ct[c]++;f=1;}
        dfs2(c,v);
        if(a[p]==1)black[v]--;
        // if(f)ct[c]--;
        dp[c]=x2;
        dp[v]=x;
        ct[c]=x3;
        ct[v]=x1;
        // if(dp[c]>=2&&ans[c])ans[v]=1;
    }
}



void solve(){

    ll n;
    cin>>n;
    fori(n+3){
        g[i].clear();

    }
    
    ll x=1;
    forn(n){
        cin>>a[i];
        if(a[i]==1){x=i;ans[i]=1;dp[i]=1;
            // black[i]=1;
        }
    }

    for(int i=0;i<n-1;i++){
        ll x,y;
        cin>>x>>y;
        g[x].pb(y);
        g[y].pb(x);
    }
    dfs(x,0);
    // forn(n){
    //     cout<<i<<" "<<dp[i]<<" "<<ct[i]<<endl;
    // }
    dfs2(x,0);
    forn(n){
        cout<<ans[i]<<" ";

    }
    cout<<endl;



}

int main(){


    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    //Redeem yourself King.
    solve();
    
   
    return 0;
}