Submission #555322

#	Time	Username	Problem	Language	Result	Execution time	Memory
555322		urosk	Mergers (JOI19_mergers)	C++14	34 / 100	353 ms	94836 KiB

mergers

// __builtin_popcount(x)
// __builtin_popcountll(x)
#define here cerr<<"===========================================\n"
#include <bits/stdc++.h>
#define ld double
#define ll long long
#define ull unsigned long long
#define llinf 100000000000000000LL // 10^17
#define iinf 2000000000 // 2*10^9
#define pb push_back
#define popb pop_back
#define fi first
#define sc second
#define endl '\n'
#define pii pair<int,int>
#define pll pair<ll,ll>
#define pld pair<ld,ld>
#define sz(a) int(a.size())
#define all(a) a.begin(),a.end()
#define ceri(a,l,r) {for(ll i_ = l;i_<=r;i_++) cerr<<a[i_]<< " ";cerr<<endl;}
using namespace std;

#define maxn 5005
ll n,k;
ll a[maxn];
ll dp[maxn][2];
ll b[maxn];
ll in[maxn];
ll out[maxn];
vector<ll> g[maxn];
vector<ll> v[maxn];
vector<ll> f[maxn];
ll ti = 1;
ll col = -1;
ll con[maxn];
ll dsu[maxn];
ll root(ll x){
    while(x!=dsu[x]){
        dsu[x] = dsu[dsu[x]];
        x = dsu[x];
    }
    return x;
}
void upd(ll x,ll y){
    x = root(x); y = root(y);
    dsu[x] = y;
}
ll root2(ll x){
    while(x!=con[x]){
        con[x] = con[con[x]];
        x = con[x];
    }
    return x;
}
void upd2(ll x,ll y){
    x = root2(x); y = root2(y);
    con[x] = y;
}
void updg(ll x,ll y){
    upd2(x,y);
    f[x].pb(y);
    f[y].pb(x);
}
vector<vector<ll> > cnt;
void dfs(ll u,ll par){
    cnt[u][b[u]] = 1;
    for(ll s : g[u]){
        if(s==par) continue;
        dfs(s,u);
        for(ll i = 1;i<=k;i++) cnt[u][i]+=cnt[s][i];
    }
    bool flag =0;
    for(ll i = 1;i<=k;i++){
        if(cnt[u][i]!=sz(v[i])&&cnt[u][i]!=0) flag = 1;
    }
    if(flag) upd(u,par);
}
ll dis[maxn];
ll getnaj(ll u,ll par){
    ll y = u;
    for(ll s : f[u]){
        if(s==par) continue;
        dis[s] = dis[u] + 1;
        ll e = getnaj(s,u);
        if(dis[e]>dis[y]) y = e;
    }
    return y;
}
ll sizi = 0;
void brisi(ll u,ll par,ll x,vector<ll> w){
    w.pb(u);
    if(u==x){
        for(ll i = 0;i<sz(w)-1;i++){
            upd(w[i],w[i+1]);
        }
        sizi-=(sz(w)-1);
        return;
    }
    for(ll s : f[u]){
        if(s==par) continue;
        brisi(s,u,x,w);
    }
}
bool vis[maxn];
int main(){
	ios_base::sync_with_stdio(false);cerr.tie(0);cout.tie(0);cin.tie(0);
    cin >> n >> k;
    iota(dsu+1,dsu+n+1,1);
    for(ll i = 1;i<=n-1;i++){
        ll x,y; cin >> x >> y;
        g[x].pb(y);
        g[y].pb(x);
    }
    cnt.resize(n+1);
    for(ll i = 0;i<=n;i++) cnt[i].resize(k+1);
    for(ll i = 1;i<=n;i++) cin >> b[i];
    for(ll i = 1;i<=n;i++){
        v[b[i]].pb(i);
    }
    dfs(1,1);
    for(ll i = 1;i<=n;i++){
        if(vis[root(i)]) continue;
        vis[root(i)] = 1;
        sizi++;
    }
    iota(con+1,con+n+1,1);
    for(ll i = 1;i<=n;i++){
        ll x = root(i);
        for(ll j : g[i]){
            ll y = root(j);
            if(root2(x)==root2(y)) continue;
            updg(x,y);
        }
    }
    //for(ll i= 1;i<=n;i++) cerr<<root(i)<< " ";
    //cerr<<endl;
    ll ans = 0;
    while(sizi>1){
        /*
        here;
        for(ll i = 1;i<=n;i++){
            if(sz(f[i])==0) continue;
            cerr<<"i: "<<i<<endl;
            for(ll x : f[i]) cerr<<x<< " ";
            cerr<<endl;
        }
        */
        ll poc = 1;
        for(ll i = 1;i<=n;i++) if(sz(f[i])>0) poc = i;
        //cerr<<"poc: "<<poc<<endl;
        fill(dis,dis+n+1,0);
        dis[0] = llinf;
        ll x = getnaj(poc,poc);
        //cerr<<x<<endl;

        fill(dis,dis+n+1,0);
        dis[0] = llinf;
        ll y = getnaj(x,x);
        brisi(x,x,y,{});
        //cerr<<y<<endl;
        for(ll i = 1;i<=n;i++) f[i].clear();
        iota(con+1,con+n+1,1);
        for(ll i = 1;i<=n;i++){
            ll x = root(i);
            for(ll j : g[i]){
                ll y = root(j);
                if(root2(x)==root2(y)) continue;
                updg(x,y);
            }
        }
        ans++;
    }
    cout<<ans<<endl;
	return 0;
}
/*
5 4
1 2
2 3
3 4
3 5
1
2
1
3
4
*/

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 0 / 14
• Subtask 4: 0 / 22
• Subtask 5: 0 / 30