Submission #829450

#	Time	Username	Problem	Language	Result	Execution time	Memory
829450		Cookie	Capital City (JOI20_capital_city)	C++14	100 / 100	507 ms	32956 KiB

capital_city

#include<bits/stdc++.h>
#include<fstream>
#pragma GCC optimize("Ofast,O3,unroll-loops")
#pragma GCC target("avx2")
using namespace std;
//ifstream fin("FEEDING.INP");
//ofstream fout("FEEDING.OUT");
#define sz(a) (int)a.size()
#define ll long long
#define pb push_back
#define forr(i, a, b) for(int i = a; i < b; i++)
#define dorr(i, a, b) for(int i = a; i >= b; i--)
#define ld long double
#define vt vector
#include<fstream>
#define fi first
#define se second
#define pll pair<ll, ll>
#define pii pair<int, int>

const ll mod = 1e9 + 7;
const int mxn = 2e5 + 5, mxr = 25e3, sq = 500, mxv = 1e6 + 5, mxvv = 130, pr = 31, inf = 1e9;
int n, k;
int sz[mxn + 1];
vt<int>adj[mxn + 1], tocol[mxn + 1];
int pa[mxn + 1], a[mxn + 1], cnt[mxn + 1], ans = inf;
bool vis[mxn + 1], vis2[mxn + 1];
vt<int>comp;
int dfs(int s, int pre){
    
    sz[s] = 1;
    for(auto i: adj[s]){
        if(i != pre && !vis[i]){
            sz[s] += dfs(i, s);
        }
    }
    return(sz[s]);
}
int centroid(int s, int pre, int need){
    for(auto i: adj[s]){
        if(i != pre && !vis[i]){
            if(sz[i] * 2 > need)return(centroid(i, s, need));
        }
    }
    return(s);
}
void go(int s, int pre){
    pa[s] = pre; comp.pb(s);
    for(auto i: adj[s]){
        if(i != pre && !vis[i]){
            go(i, s);
        }
    }
}
void solve(int s){
    comp.clear();
    int c = centroid(s, -1, dfs(s, -1));
    go(c, -1);
   
    vis[c] = 1;
    
    for(auto i: comp){
        tocol[a[i]].pb(i);
    }
    queue<int>q;
    vis2[a[c]] = 1;
    q.push(a[c]);
    int need = 1; bool ok = 1;
    while(!q.empty()){
        int nw = q.front(); q.pop();
        if(cnt[nw] != sz(tocol[nw])){
            ok = 0; break;
        }
        for(auto i: tocol[nw]){
            int to = pa[i];
            if(to != -1){
                if(!vis2[a[to]]){
                    vis2[a[to]] = 1;  need++;
                    q.push(a[to]);
                
                }
            }
        }
    }
    for(auto i: comp){
        tocol[a[i]].clear(); vis2[a[i]] = 0; 
    }
    if(ok)ans = min(ans, need - 1);
    for(auto i: adj[c]){
        if(!vis[i]){
            solve(i);
        }
    }
    
}
int main(){
     ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n >> k;
    for(int i = 1; i < n; i++){
        int u, v; cin >> u >> v;
        adj[u].pb(v); adj[v].pb(u);
    }
    for(int i = 1; i <= n; i++){
        cin >> a[i]; cnt[a[i]]++;
    }
    solve(1);
    cout << ans;
    return(0);
}

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