Submission #1174584

#	Time	Username	Problem	Language	Result	Execution time	Memory
1174584		hengliao	Balanced Tree (info1cup18_balancedtree)	C++20	0 / 100	143 ms	13916 KiB

balancedtree

#include<bits/stdc++.h>
using namespace std;

#define F first
#define S second
#define pb push_back
#define vll vector<ll>
#define pll pair<ll, ll>

typedef int ll;

const ll mxN=5e5+5;
const ll inf=1e9;

ll n;
vll adj[mxN];
array<array<array<ll, 2> ,2>, 2> dp[mxN];
array<bool, 2> dp2[mxN];
ll c[mxN];
ll tar;
ll sz[mxN];

void dfs(ll cur, ll par){
    for(ll i=0;i<2;i++){
        for(ll j=0;j<2;j++){
            for(ll k=0;k<2;k++){
                dp[cur][i][j][k]=inf;
            }
        }
    }
    sz[cur]=1;
    // if(c[cur]==-1){
    //     dp[cur][0][1][0]=0;
    //     dp[cur][1][1][0]=0;
    // }
    // else{
    //     dp[cur][c[cur]][1][0]=0;
    // }
    if(c[cur]!=-1){
        dp2[cur][c[cur]^1]=0;
    }
    for(auto &chd:adj[cur]){
        if(chd==par) continue;
        dfs(chd, cur);
        
        array<array<array<ll, 2> ,2>, 2> new_dp;
        for(ll i=0;i<2;i++){
            for(ll j=0;j<2;j++){
                for(ll k=0;k<2;k++){
                    new_dp[i][j][k]=inf;
                }
            }
        }
        
        for(ll i=0;i<2;i++){
            if(dp2[chd][i]){
                for(ll j=0;j<2;j++){
                    for(ll k=0;k<2;k++){
                        for(ll t=0;t<2;t++){
                            if(dp[cur][j][k][t]==inf) continue;
                            if(i==j){
                                new_dp[j][0][t]=min(new_dp[j][0][t], dp[cur][j][k][t]);
                            }
                            else{
                                new_dp[j][k][0]=min(new_dp[j][k][0], 1);
                            }
                        }
                    }
                }
            }
        }
        for(ll i=0;i<2;i++){
            if(dp2[cur][i]){
                for(ll j=0;j<2;j++){
                    for(ll k=0;k<2;k++){
                        for(ll t=0;t<2;t++){
                            if(dp[chd][j][k][t]==inf) continue;
                            if(i==j){
                                new_dp[i][0][t]=min(new_dp[i][0][t], dp[chd][j][k][t]+1);
                            }
                            else{
                                new_dp[i][0][k]=min(new_dp[i][0][k], 1);
                            }
                        }
                    }
                }
            }
        }
        for(ll i=0;i<2;i++){
            for(ll j=0;j<2;j++){
                for(ll k=0;k<2;k++){
                    if(dp[cur][i][j][k]==inf) continue;
                    for(ll t=0;t<2;t++){
                        for(ll x=0;x<2;x++){
                            for(ll y=0;y<2;y++){
                                if(dp[chd][t][x][y]==inf) continue;
                                if(i==t){
                                    if(k==0 && y==0){
                                        new_dp[i][0][0]=min(min(new_dp[i][0][0], dp[cur][i][j][k]), dp[chd][t][x][y]+1);
                                    }
                                    else if(k==0){
                                        if(dp[cur][i][j][k]+dp[chd][t][x][y]+1<=tar){
                                            new_dp[i][0][0]=min(min(new_dp[i][0][0], dp[chd][t][x][y]+1), dp[cur][i][j][k]);
                                        }
                                        else{
                                            new_dp[i][0][1]=min(new_dp[i][0][1], dp[chd][t][x][y]+1);
                                        }
                                    }
                                    else if(y==0){
                                        if(dp[cur][i][j][k]+dp[chd][t][x][y]+1<=tar){
                                            new_dp[i][0][0]=min(min(new_dp[i][0][0], dp[cur][i][j][k]), dp[chd][t][x][y]+1);
                                        }
                                        else{
                                            new_dp[i][0][1]=min(new_dp[i][0][1], dp[cur][i][j][k]);
                                        }
                                    }
                                    else{
                                        if(dp[cur][i][j][k]+dp[chd][t][x][y]+1<=tar){
                                            new_dp[i][0][0]=min(min(new_dp[i][0][0], dp[cur][i][j][k]), dp[chd][t][x][y]+1);
                                        }
                                        else{
                                            new_dp[i][0][1]=min(min(new_dp[i][0][1], dp[cur][i][j][k]), dp[chd][t][x][y]+1);
                                        }
                                    }
                                }
                                else{
                                    // if(j==0){
                                        new_dp[i][0][0]=1;
                                    // }
                                    // else if(x==0){
                                    //     if(dp[cur][i][j][k]+dp[chd][t][x][y]+1<=tar){
                                    //         new_dp[i][0][0]=min(new_dp[i][0][0], dp[chd][t][x][y]+1);
                                    //     }
                                    //     else{
                                    //         new_dp[i][0][1]=min(new_dp[i][0][1], dp[cur][i][j][k]);
                                    //     }
                                    // }
                                    // else{
                                    //     if(dp[cur][i][j][k]+dp[chd][t][x][y]+1<=tar){
                                    //         new_dp[i][0][0]=min(min(new_dp[i][0][0], dp[cur][i][j][k]), dp[chd][t][x][y]+1);
                                    //     }
                                    //     else{
                                    //         new_dp[i][0][1]=min(min(new_dp[i][0][1], dp[cur][i][j][k]), dp[chd][t][x][y]+1);
                                    //     }
                                    // }
                                }
                            }
                        }
                    }
                }
            }
        }
        for(ll i=0;i<2;i++){
            for(ll j=0;j<2;j++){
                if(i==j) continue;
                if(dp2[cur][i] && dp2[chd][j]){
                    if(sz[chd]==1){
                        if(sz[cur]==1)
                            new_dp[i][1][1]=1;
                        else 
                            new_dp[i][0][1]=1;
                    }
                    else{
                        if(sz[cur]==1)
                            new_dp[i][1][0]=1;
                        else
                            new_dp[i][0][0]=1;
                    }
                }
            }
        }
        dp2[cur][0]&=dp2[chd][0];
        dp2[cur][1]&=dp2[chd][1];
        for(ll i=0;i<2;i++){
            for(ll j=0;j<2;j++){
                for(ll k=0;k<2;k++){
                    if(k==1 && new_dp[i][j][k]>=tar){
                        new_dp[i][j][k]=inf;
                    }
                }
            }
        }
        sz[cur]+=sz[chd];
        dp[cur]=new_dp;
        // if(cur==5){
            // cout<<cur<<" after "<<chd<<'\n';
            // for(ll i=0;i<2;i++){
            //     for(ll j=0;j<2;j++){
            //         for(ll k=0;k<2;k++){
            //             cout<<i<<' '<<j<<' '<<k<<' '<<dp[cur][i][j][k]<<'\n';
            //         }
            //     }
            // }
        // }
    }
}

bool check(ll tt){
    for(ll i=1;i<=n;i++){
        for(ll j=0;j<2;j++){
            for(ll k=0;k<2;k++){
                for(ll t=0;t<2;t++){
                    dp[i][j][k][t]=inf;
                    dp2[i][j]=1;
                }
            }
        }
    }
    tar=tt;
    dfs(1, 0);
    for(ll i=0;i<2;i++){
        if(dp2[1][i]) return true;
    }
    for(ll i=0;i<2;i++){
        if(dp[1][i][0][0]!=inf) return true;
    }
    return false;
}

void solve(){
    cin>>n;
    for(ll i=1;i<=n;i++){
        adj[i].clear();
    }
    for(ll i=0;i<n-1;i++){
        ll u, v;
        cin>>u>>v;
        adj[u].pb(v);
        adj[v].pb(u);
    }
    for(ll i=1;i<=n;i++){
        cin>>c[i];
    }
    // check(2);
    // for(ll i=1;i<=n;i++){
    //     cout<<"node: "<<i<<'\n';
    //     for(ll j=0;j<2;j++){
    //         for(ll k=0;k<2;k++){
    //             for(ll t=0;t<2;t++){
    //                 cout<<j<<' '<<k<<' '<<t<<" : "<<dp[i][j][k][t]<<'\n';
    //             }
    //         }
    //     }
    // }
    
    ll lef=1, rig=n-1;
    ll ans=-1;
    while(lef<=rig){
        ll mid=(lef+rig)/2;
        if(check(mid)){
            ans=mid;
            rig=mid-1;
        }
        else{
            lef=mid+1;
        }
    }
    cout<<ans<<'\n';
    if(ans==-1) return;
    for(ll i=1;i<=n;i++){
        if(c[i]==-1 && i==4){
            cout<<1<<' ';
        }
        else if(c[i]==-1){
            cout<<0<<' ';
        }
        else{
            cout<<c[i]<<' ';
        }
    }
    cout<<'\n';
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    int t;
    cin>>t;
    while(t--){
        solve();
    }

    

    return 0;
}

• Subtask 1: 0 / 10
• Subtask 2: 0 / 13
• Subtask 3: 0 / 26
• Subtask 4: 0 / 45
• Subtask 5: 0 / 6