Submission #446426

#	Time	Username	Problem	Language	Result	Execution time	Memory
446426		MladenP	Capital City (JOI20_capital_city)	C++17	100 / 100	686 ms	39196 KiB

capital_city

#include <bits/stdc++.h>
#define PRINT(x) cerr<<#x<<'='<<x<<endl;
#define NL(x) " \n"[(x)]
#define lld long long
#define pil pair<int,lld>
#define pli pair<lld,int>
#define pll pair<lld,lld>
#define pii pair<int,int>
#define pb push_back
#define fi first
#define se second
#define mid (l+r)/2
#define endl '\n'
#define all(a) begin(a),end(a)
#define sz(a) int((a).size())
#define LINF 1000000000000000LL
#define INF 1000000000
#define EPS 1e-9
using namespace std;
#define MAXN 200010
int N, K, C[MAXN], siz[MAXN], par[MAXN], sol;
bool vis[MAXN], col[MAXN], pos[MAXN], ovde[MAXN];
vector<int> adj[MAXN], v, pC[MAXN];
int dfsSize(int node, int prev) {
    siz[node] = 1;
    for(auto x : adj[node]) {
        if(!vis[x] && x != prev) siz[node] += dfsSize(x, node);
    }
    return siz[node];
}
int dfsCent(int node, int prev, int S) {
    for(auto x : adj[node]) {
        if(!vis[x] && x != prev && siz[x] > S/2) return dfsCent(x, node, S);
    }
    return node;
}
void dfsParent(int node, int prev) {
    par[node] = prev;
    ovde[node] = true;
    v.pb(node);
    for(auto x : adj[node]) {
        if(!vis[x] && x != prev) dfsParent(x, node);
    }
}
vector<int> nv;
int boja = 0;
bool addColor(int color) {
    if(col[color]) return true;
    boja++;
    col[color] = true;
    for(auto x : pC[color]) {
        if(!ovde[x]) return false;
        if(!pos[x]) nv.pb(x);
    }
    return true;
}
void centDecomp(int node) {
    int S = dfsSize(node, node);
    node = dfsCent(node, node, S);
    vis[node] = 1;
    dfsParent(node, node);
    boja = 0;
    if(!addColor(C[node])) {
        nv.clear();
        boja = K;
    }
    while(!nv.empty()) {
        int cur = nv.back(); nv.pop_back();
        if(pos[cur]) continue;
        while(1) {
            pos[cur] = true;
            if(!addColor(C[cur])) {
                nv.clear();
                boja = K;
                break;
            }
            if(par[cur] == cur || pos[par[cur]]) break;
            cur = par[cur];
        }
    }
    sol = min(sol, boja);
    for(auto x : v) pos[x] = ovde[x] = col[C[x]] = 0;
    v.clear();
    for(auto x : adj[node]) if(!vis[x]) centDecomp(x);
}
int main() {
    ios::sync_with_stdio(false); cin.tie(0);
    cin >> N >> K;
    for(int i = 1; i <= N-1; i++) {
        int x, y; cin >> x >> y;
        adj[x].pb(y);
        adj[y].pb(x);
    }
    for(int i = 1; i <= N; i++) {
        cin >> C[i];
        pC[C[i]].pb(i);
    }
    sol = K;
    centDecomp(1);
    cout << sol-1;
    return 0;
}

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