Submission #612558

#	Time	Username	Problem	Language	Result	Execution time	Memory
612558		alireza_kaviani	Capital City (JOI20_capital_city)	C++17	11 / 100	2051 ms	524288 KiB

capital_city

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

typedef long long ll;
typedef pair<ll, ll> pll;
typedef pair<int, int> pii;

#define all(x)      (x).begin(),(x).end()
#define X           first
#define Y           second
#define sep         ' '
#define endl        '\n'
#define SZ(x)       ll(x.size())

const ll MAXN = 2e5 + 10;
const ll LOG = 20;
const ll MAX = MAXN * LOG + MAXN;
const ll INF = 8e18;
const ll MOD = 1e9 + 7; //998244353; //1e9 + 9;

int n , k , cnt , ans = MOD , C[MAXN] , par[LOG][MAXN] , H[MAXN] , mark[MAX];
vector<int> tree[MAXN] , adj[MAX] , radj[MAX] , vec[MAXN] , ord;

void addEdge(int v , int u){
    adj[v].push_back(u);
    radj[u].push_back(v);
}

void PreDFS(int v , int p = 0){
    par[0][v] = p;
    H[v] = H[p] + 1;
    for(int u : tree[v]){
        if(u == p)  continue;
        PreDFS(u , v);
    }
}

int getPar(int v , int h){
    for(int i = 0 ; i < LOG ; i++){
        if(h & (1 << i)){
            v = par[i][v];
        }
    }
    return v;
}

int LCA(int v , int u){
    if(H[v] > H[u]) swap(v , u);
    u = getPar(u , H[u] - H[v]);
    if(v == u)  return v;
    for(int i = LOG - 1 ; i >= 0 ; i--){
        if(par[i][v] != par[i][u]){
            v = par[i][v];
            u = par[i][u];
        }
    }
    return par[0][v];
}

void add(int x , int v , int h){
    for(int i = 0 ; i < LOG ; i++){
        if(h & (1 << i)){
            addEdge(x , i * MAXN + v);
            v = par[i][v];
        }
    }
}

void DFS(int v){
    mark[v] = 1;
    for(int u : radj[v]){
        if(!mark[u]){
            DFS(u);
        }
    }
    ord.push_back(v);
}

void DFS2(int v){
    if(LOG * MAXN + 1 <= v && v <= LOG * MAXN + k){
        cnt++;
    }
    mark[v] = 1;
    for(int u : adj[v]){
        if(!mark[u]){
            DFS2(u);
        }
    }
}

void DFS3(int v){
    mark[v] = 2;
    for(int u : radj[v]){
        if(mark[u] != 2){
            DFS3(u);
        }
    }
}

void SCC(){
    for(int i = 0 ; i < MAX ; i++){
        if(!mark[i]){
            DFS(i);
        }
    }
    fill(mark , mark + MAX , 0);
    reverse(all(ord));
    for(int i : ord){
        if(mark[i]) continue;
        cnt = 0;
        DFS2(i);
        DFS3(i);
        if(cnt > 0){
            ans = min(ans , cnt);
        }
    }
}

int main() {
    ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

    cin >> n >> k;
    for(int i = 1 ; i < n ; i++){
        int v , u;
        cin >> v >> u;
        tree[v].push_back(u);
        tree[u].push_back(v);
    }
    for(int i = 1 ; i <= n ; i++){
        cin >> C[i];
        vec[C[i]].push_back(i);
    }
    PreDFS(1);

    for(int i = 1 ; i < LOG ; i++){
        for(int j = 1 ; j <= n ; j++){
            par[i][j] = par[i - 1][par[i - 1][j]];
            addEdge(i * MAXN + j , (i - 1) * MAXN + j);
            addEdge(i * MAXN + j , (i - 1) * MAXN + par[i - 1][j]);
        }
    }
    for(int i = 1 ; i <= n ; i++){
        addEdge(i , LOG * MAXN + C[i]);
    }

    for(int i = 1 ; i <= k ; i++){
        for(int j = 0 ; j + 1 < SZ(vec[i]) ; j++){
            int v = vec[i][j] , u = vec[i][j + 1];
            int lca = LCA(v , u);
            add(LOG * MAXN + i , v , H[v] - H[lca]);
            add(LOG * MAXN + i , u , H[u] - H[lca]);
            addEdge(LOG * MAXN + i , lca);
        }
    }

    SCC();
    cout << ans - 1 << endl;

    return 0;
}
/*

*/

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