Submission #230697

#	Time	Username	Problem	Language	Result	Execution time	Memory
230697		WLZ	Capital City (JOI20_capital_city)	C++14	100 / 100	1419 ms	51992 KiB

capital_city

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

const int INF = 0x3f3f3f3f;

int k, ans = INF;
vector< vector<int> > g, towns;
vector<int> c, sz, par, cur;
vector<bool> used;

void dfs(int u, int p, vector<int> &out) {
    sz[u] = 1; par[u] = p;
    out.push_back(u);
    for (auto& v : g[u]) {
        if (v != p && !used[v]) {
            dfs(v, u, out);
            sz[u] += sz[v];
        }
    }
}

int find_centroid(int u, int p, int n) {
    for (auto& v : g[u]) {
        if (v != p && sz[v] > n / 2 && !used[v]) {
            return find_centroid(v, u, n);
        }
    }
    return u;
}

void solve(int u) {
    vector<int> junk;
    dfs(u, -1, junk);
    int centroid = find_centroid(u, -1, sz[u]);
    dfs(centroid, -1, cur);
    map<int, int> cnt;
    for (auto& x : cur) {
        cnt[c[x]]++;
    }
    set<int> st;
    if (cnt[c[centroid]] == (int) towns[c[centroid]].size()) {
        bool pos = true;
        queue<int> q;
        for (auto& x : towns[c[centroid]]) {
            q.push(x);
        }
        st.insert(c[centroid]);
        while (!q.empty()) {
            int x = q.front(); q.pop();
            if (x != centroid && !st.count(c[par[x]])) {
                if (cnt[c[par[x]]] != (int) towns[c[par[x]]].size()) {
                    pos = false;
                    break;
                }
                for (auto& y : towns[c[par[x]]]) {
                    q.push(y);
                }
                st.insert(c[par[x]]);
            }
        }
        if (pos) {
            ans = min(ans, (int) st.size());
        }
    }
    used[centroid] = true;
    cur.clear();
    for (auto& v : g[centroid]) {
        if (!used[v]) {
            solve(v);
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int n;
    cin >> n >> k;
    g.resize(n + 1);
    for (int i = 0; i < n - 1; i++) {
        int a, b;
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    c.resize(n + 1);
    towns.resize(k + 1);
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
        towns[c[i]].push_back(i);
    }
    par.resize(n + 1);
    sz.resize(n + 1);
    used.assign(n + 1, false);
    solve(1);
    cout << ans - 1 << '\n';
    return 0;
}

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