Submission #1262090

#	Time	Username	Problem	Language	Result	Execution time	Memory
1262090		nguyn	수도 (JOI20_capital_city)	C++20	100 / 100	549 ms	40084 KiB

capital_city

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

#define ll long long
#define F first
#define S second
#define pb push_back
#define pii pair<int, int>

const int N = 2e5 + 5;

int n, k;
vector<int> g[N];
int del[N];
int sz[N];
int par[N];
int c[N];
int vis[N];
int cnt_col[N];
vector<int> col[N];
vector<int> cand;
vector<int> vec;
int res = 1e9;

void pre_dfs(int u, int p) {
    sz[u] = 1;
    for (int v : g[u]) {
        if (v == p || del[v]) continue;
        pre_dfs(v, u);
        sz[u] += sz[v];
    }
}

int find_centroid(int u, int p, int n) {
    for (int v : g[u]) {
        if (v == p || del[v]) continue;
        if (sz[v] > n / 2) return find_centroid(v, u, n);
    }
    return u;
}

void dfs(int u, int p, int root) {
    par[u] = p;
    if (c[u] == c[root]) cand.pb(u);
    vec.pb(u);
    cnt_col[c[u]]++;
    for (int v : g[u]) {
        if (v == p || del[v]) continue;
        dfs(v, u, root);
    }
}

void solve(int u) {
    pre_dfs(u, 0);
    int n = sz[u];
    int root = find_centroid(u, 0, n);
    dfs(root, 0, root);
    del[root] = 1;

    for (int v : cand) {
        vis[v] = 1;
    }
    if (cnt_col[c[root]] == col[c[root]].size()) {
        int ans = 0;
        while(!cand.empty()) {
            int cur = cand.back();
            cand.pop_back();
            int v = par[cur];
            while(!vis[v] && v != 0) {
                if (cnt_col[c[v]] != col[c[v]].size()) {
                    ans = 1e9;
                    break;
                }
                for (int x : col[c[v]]) {
                    cand.pb(x);
                    vis[x] = 1;
                }
                ans++;
                v = par[v];
            }
            if (ans == 1e9) break;
        }
//        cout << root << ' ' << ans << ":\n";
//        for (int x : vec) {
//            cout << x << ' ';
//        } cout << '\n';
        res = min(res, ans);
    }
//    cout << root << ":\n";
    for (int v : vec) {
//        cout << v << ' ';
        vis[v] = 0;
        cnt_col[c[v]] = 0;
    }
//    cout << '\n';
    cand.clear();
    vec.clear();

    for (int v : g[root]) {
        if (del[v]) continue;
        solve(v);
    }
}

signed main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        g[u].pb(v);
        g[v].pb(u);
    }
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
        col[c[i]].pb(i);
    }
    solve(1);
    cout << res;
}

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