Submission #1156007

#	Time	Username	Problem	Language	Result	Execution time	Memory
1156007		Pekiban	Capital City (JOI20_capital_city)	C++20	100 / 100	578 ms	38112 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;
#define pb push_back

const int N = 2e5 + 5;
vector<int> g[N], h[N];
int sz[N], c[N], p[N], f[N];
bool vis[N], visC[N], bad[N];

void dfssz(int s, int e = -1) {
    sz[s] = 1;
    for (auto u : g[s]) {
        if (u == e || vis[u]) continue;
        dfssz(u, s);
        sz[s] += sz[u];
    }
}
int centroid(int s, int n, int e = -1) {
    for (auto u : g[s]) {
        if (u == e || vis[u])   continue;
        if (sz[u] > n/2)    return centroid(u, n, s);
    }
    return s;
}
void dfsroot(int s, int e = -1) {
    for (auto u : g[s]) {
        if (u == e || vis[u])   continue;
        p[u] = s;
        dfsroot(u, s);
    }
}
void upd(int s, int e, int x) {
    if (f[c[s]] != x && f[c[s]] && x)   bad[c[s]] = 1;
    f[c[s]] = x;
    for (auto u : g[s]) {
        if (vis[u] || u == e)   continue;
        upd(u, s, x);
    }
}
int ans = 1e9;
void solve(int R, int n) {
    dfssz(R);
    int nd = centroid(R, n);
    dfssz(nd);
    if (!bad[c[nd]]) {
        dfsroot(nd);
        queue<int> Q; Q.push(c[nd]);
        vector<int> cl = {c[nd]};
        visC[c[nd]] = 1;
        bool ok = 1;
        while (!Q.empty() && ok) {
            int C = Q.front(); Q.pop();
            for (auto u : h[C]) {
                if (!ok)    break;
                if (u == nd)    continue;
                int v = p[u];
                while (!visC[c[v]]) {
                    if (bad[c[v]]) {
                        ok = 0;
                        break;
                    }
                    visC[c[v]] = 1;
                    Q.push(c[v]);
                    cl.pb(c[v]);
                    v = p[v];
                }
            }
        }
        bad[c[nd]] = 1;
        for (auto i : cl)   visC[i] = 0;
        if (ok) ans = min(ans, (int)cl.size() - 1);
    }
    for (auto u : g[nd]) {
        if (vis[u])    continue;
        upd(u, nd, u);
    }
    upd(nd, -1, 0);
    vis[nd] = 1;
    for (auto u : g[nd]) {
        if (vis[u]) continue;
        solve(u, sz[u]);
    }
}
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    int n, k;
    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];
        h[c[i]].pb(i);
    }
    solve(1, n);
    cout << ans << '\n';
}

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