Submission #582737

#	Time	Username	Problem	Language	Result	Execution time	Memory
582737		elkernos	Capital City (JOI20_capital_city)	C++17	100 / 100	1207 ms	45260 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

int main()
{
    cin.tie(nullptr)->sync_with_stdio(0);
    int n, k;
    cin >> n >> k;
    vector<int> m(n + 1);
    vector<vector<int>> gg(k + 1), g(n + 1);
    for(int i = 1; i < n; i++) {
        int a, b;
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    for(int i = 1; i <= n; i++) {
        cin >> m[i];
        gg[m[i]].push_back(i);
    }
    vector<int> sub(n + 1), par(n + 1), cnt(k + 1);
    vector<bool> used(n + 1), onq(k + 1);
    function<int(int, int, int)> subsz = [&](int u, int p, int c) {
        sub[u] = 1;
        par[u] = p;
        cnt[m[u]] += c;
        onq[m[u]] = 0;
        for(int to : g[u]) {
            if(to == p || used[to]) {
                continue;
            }
            sub[u] += subsz(to, u, c);
        }
        return sub[u];
    };
    function<int(int, int, int)> getcen = [&](int u, int p, int d) {
        for(int to : g[u]) {
            if(to == p || used[to]) {
                continue;
            }
            if(sub[to] > d / 2) {
                return getcen(to, u, d);
            }
        }
        return u;
    };
    int ans = n;
    function<void(int)> rek = [&](int u) {
        subsz(u, u, 0);
        u = getcen(u, u, sub[u]);
        subsz(u, u, 1);
        int use = 0;
        queue<int> q;
        onq[m[u]] = 1;
        q.push(m[u]);
        while(!q.empty()) {
            int v = q.front();
            q.pop();
            if(cnt[v] != (int)gg[v].size()) {
                use = n;
                break;
            }
            for(int to : gg[v]) {
                int look = m[par[to]];
                if(onq[look] == 0) {
                    q.push(look);
                    onq[look] = 1;
                    use++;
                }
            }
        }
        ans = min(ans, use);
        subsz(u, u, -1);
        used[u] = true;
        for(int to : g[u]) {
            if(used[to]) {
                continue;
            }
            rek(to);
        }
    };
    rek(1);
    cout << ans << '\n';
}

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