Submission #359152

#	Time	Username	Problem	Language	Result	Execution time	Memory
359152		Nima_Naderi	Capital City (JOI20_capital_city)	C++14	100 / 100	1157 ms	52708 KiB

capital_city

///In the name of GOD
//#pragma GCC optimize("O2")
#include<bits/stdc++.h>
using namespace std;

typedef long long ll;
const ll MXN = 4e5 + 10;
ll n, k, ans = 1e9;
ll A[MXN], sub[MXN], dad[MXN], Cnt[MXN], All[MXN];
vector<ll> adj[MXN], Ver[MXN], vec, col;
bool hide[MXN], bad[MXN], mark[MXN], vis[MXN];
void plant(ll u, ll par){
    sub[u] = 1;
    for(auto v : adj[u]){
        if(v != par && !hide[v]){
            plant(v, u), sub[u] += sub[v];
        }
    }
}
void Pre(ll u, ll par, ll dt){
    dad[u] = par, mark[u] = 0;
    if(dt == +1 && Cnt[A[u]] == 0) col.push_back(A[u]);
    if(dt == +1) Cnt[A[u]] ++;
    else         Cnt[A[u]] = 0;
    for(auto v : adj[u]){
        if(v != par && !hide[v]) Pre(v, u, dt);
    }
}
ll CRD(ll u, ll par, ll val){
    for(auto v : adj[u]){
        if(v == par || hide[v]) continue;
        if(sub[v] >= val) return CRD(v, u, val);
    }
    return u;
}
void DMS(ll u, ll h){
    plant(u, -1);
    ll cent = CRD(u, -1, (sub[u] + 1) / 2);//attention
    col.clear(), vec.clear();
    Pre(cent, -1, +1);
    for(auto c : col){
        if(Cnt[c] != All[c]) bad[c] = 1;
        else                 bad[c] = 0;
        vis[c] = 0;
    }
    if(!bad[A[cent]]){
        vec.push_back(cent);
        ll tot = 0, pnt = 0; bool ok = 1;
        while(pnt < vec.size()){
            ll now = vec[pnt]; pnt ++;
            if(mark[now] || now == -1) continue;
            ll vr = now;
            while(vr != -1 && !mark[vr]){
                mark[vr] = 1;
                if(!vis[A[vr]]){
                    tot ++; vis[A[vr]] = 1;
                    ok &= (!bad[A[vr]]);
                    if(!ok) break;
                    for(auto X : Ver[A[vr]]){
                        if(X == vr){
                            assert(mark[X]);
                            continue;
                        }
                        assert(!mark[X]);
                        vec.push_back(X);
                    }
                }
                if(!ok) break;
                vr = dad[vr];
            }
            if(!ok) break;
        }
        for(auto c : col) vis[c] = 0;
        if(ok) ans = min(ans, tot);
    }
    col.clear(), Pre(cent, -1, -1);
    hide[cent] = 1;
    for(auto v : adj[cent]){
        if(!hide[v]) DMS(v, h + 1);
    }
}
int main(){
    ios::sync_with_stdio(0);cin.tie(0); cout.tie(0);
    cin >> n >> k;
    for(int i = 1; i < n; i ++){
        ll u, v; cin >> u >> v;
        adj[u].push_back(v), adj[v].push_back(u);
    }
    for(int i = 1; i <= n; i ++){
        cin >> A[i], All[A[i]] ++;
        Ver[A[i]].push_back(i);
    }
    DMS(1, 0);
    cout << ans - 1 << '\n';
    return 0;
}
/*!
    HE'S AN INSTIGATOR,
    ENEMY ELIMINATOR,
    AND WHEN HE KNOCKS YOU BETTER
    YOU BETTER LET HIM IN.
*/
//! N.N

Compilation message (stderr)

capital_city.cpp: In function 'void DMS(ll, ll)':
capital_city.cpp:49:19: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   49 |         while(pnt < vec.size()){
      |               ~~~~^~~~~~~~~~~~

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59