Submission #1154175

#	Time	Username	Problem	Language	Result	Execution time	Memory
1154175		pokmui9909	Capital City (JOI20_capital_city)	C++17	1 / 100	481 ms	46972 KiB

capital_city

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

#define x first
#define y second

ll N, K, A[200005], Tot[200005], S[200005], vis[200005];
vector<ll> T[200005];

ll Size(ll u, ll p){
    S[u] = 1;
    for(auto v : T[u]){
        if(v != p && !vis[v]) S[u] += Size(v, u);
    }
    return S[u];
}
ll Cent(ll u, ll p, ll n){
    for(auto v : T[u]){
        if(v != p && !vis[v] && S[v] > n / 2) return Cent(v, u, n);
    }
    return u;
}

ll Par[200005], chk[200005], Ans = 1e18;
vector<ll> U[200005];

void dfs(ll u, ll p, vector<ll> &V, vector<ll> &Del){
    Par[u] = p;
    V.push_back(A[u]); Del.push_back(u);
    U[A[u]].push_back(u);
    for(auto v : T[u]){
        if(v == p || vis[v]) continue;
        dfs(v, u, V, Del);
    }
}

void dnc(ll _u){
    ll u = Cent(_u, -1, Size(_u, -1));
    vis[u] = 1;
    vector<ll> V, Del;
    dfs(u, u, V, Del);
    sort(V.begin(), V.end());
    V.erase(unique(V.begin(), V.end()), V.end());
    ll ok = 1;
    for(auto e : V){
        if(U[e].size() != Tot[e]){
            ok = 0; break;
        }
    }
    if(ok){
        queue<ll> Q;
        ll c = 0;
        for(auto i : U[A[u]]){
            chk[i] = 1; Q.push(i);
        }
        while(!Q.empty()){
            ll k = Q.front(); Q.pop();
            if(k == u) continue;
            k = Par[k];
            if(chk[k] || U[A[k]].empty()) continue;
            for(auto i : U[A[k]]){
                chk[i] = 1; Q.push(i);
                U[A[i]].clear();
            }
            U[A[k]].clear();
            c++;
            chk[k] = 1;
            Q.push(k);
        }
        Ans = min(Ans, c);
    }
    for(auto i : Del){
        chk[i] = 0; U[A[i]].clear();
    }
    for(auto v : T[u]){
        if(!vis[v]) dnc(v);
    }
}

int main(){
    cin.tie(0) -> sync_with_stdio(0);

    cin >> N >> K;
    for(ll i = 1; i < N; i++){
        ll u, v; cin >> u >> v;
        T[u].push_back(v); T[v].push_back(u);
    }
    for(ll i = 1; i <= N; i++){
        cin >> A[i];
        Tot[A[i]]++;
    }
    dnc(1);
    cout << Ans;
}

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