Submission #921464

#	Time	Username	Problem	Language	Result	Execution time	Memory
921464		lolismek	Capital City (JOI20_capital_city)	C++14	1 / 100	224 ms	41720 KiB

capital_city

#include <algorithm>
#include <iostream>
#include <climits>
#include <vector>
#include <deque>
#include <queue>
#include <stack>
#include <map>
#include <set>

#include <iomanip>
#include <cassert>

#include <random>
#include <chrono>

#include <unordered_set>
#include <unordered_map>

// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")

using ull = unsigned long long;
using ll = long long;

//#define int __int128
//#define int ll
#define pii pair <int, int>
#define all(a) (a).begin(), (a).end()
#define fr first
#define sc second
#define pb push_back
#define lb lower_bound
#define ub upper_bound

#define vt vector
#define FOR(a, b) for(int i = (a); i <= (b); i++)
#define FORr(a, b) for(int i = (a); i >= (b); i--)
#define sz(x) (int)(x).size()

#define YES cout << "YES\n"
#define NO cout << "NO\n"

using namespace std;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

int rangerng(int l, int r){
    return uniform_int_distribution<>(l, r)(rng);
}

////////////////////////////////////////////////////////////////////////////////////

const int NMAX = 2e5;

int c[NMAX + 1];
vt <int> adj[NMAX + 1];

int sz[NMAX + 1];
bool proc[NMAX + 1];

int f[NMAX + 1];

void recalc_sz(int node, int par){
    sz[node] = 1;

    for(int child : adj[node]){
        if(!proc[child] &&child != par){
            recalc_sz(child, node);
            sz[node] += sz[child];
        }
    }
}

int get_centroid(int node, int par, int lim){
    for(int child : adj[node]){
        if(!proc[child] && child != par && sz[child] > lim){
            return get_centroid(child, node, lim);
        }
    }
    return node;
}

int parent[NMAX + 1];
vt <int> inds[NMAX + 1];

void dfs(int node, int par){
    parent[node] = par;
    inds[c[node]].clear();
    for(int vec : adj[node]){
        if(!proc[vec] && vec != par){
            dfs(vec, node);
        }
    }
}

void dfs2(int node, int par){
    inds[c[node]].pb(node);
    for(int vec : adj[node]){
        if(!proc[vec] && vec != par){
            dfs2(vec, node);
        }
    }
}

int mini = 1e9;
void decomp(int node){
    recalc_sz(node, 0);
    int centroid = get_centroid(node, 0, sz[node] / 2);
    proc[centroid] = true;

    dfs(centroid, 0);
    dfs2(centroid, 0);

    queue <int> Q;
    unordered_set <int> S;
    Q.push(c[centroid]);
    S.insert(c[centroid]);

    int adds = 1;
    bool ok = 1;
    while(!Q.empty()){
        int col = Q.front();
        Q.pop();

        if(f[col] != sz(inds[col])){
            ok = 0;
        }

        for(int node : inds[col]){
            if(node != centroid && S.find(c[parent[node]]) == S.end()){
                S.insert(c[parent[node]]);
                Q.push(c[parent[node]]);
                adds++;
            }
        }
    }

    if(ok){
        mini = min(mini, adds - 1);
    }
}   

void solve(){
    int n, k;
    cin >> n >> k;

    FOR(1, n - 1){
        int a, b;
        cin >> a >> b;
        adj[a].pb(b);
        adj[b].pb(a);
    }

    FOR(1, n){
        cin >> c[i];
        f[c[i]]++;
    }    

    decomp(1);

    cout << mini << endl;
}   


signed main(){

    //ios_base::sync_with_stdio(false);
    //cin.tie(0);

    int T;
    //cin >> T;
    T = 1;

    while(T--){
        solve();
    }

    return 0;
}

/*
12 4
7 9
1 3
4 6
2 4
10 12
1 2
2 10
11 1
2 8
5 3
6 7
3
1
1
2
4
3
3
2
2
3
4
4
*/

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