Submission #332606

#	Time	Username	Problem	Language	Result	Execution time	Memory
332606		12tqian	Capital City (JOI20_capital_city)	C++17	100 / 100	1582 ms	45548 KiB

capital_city

#include <bits/stdc++.h>

int main() {
    using namespace std;
    ios_base::sync_with_stdio(0);
    int n, k;
    cin >> n >> k;
    vector<vector<int>> adj(n);
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        --u, --v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    vector<int> c(n);
    vector<int> count(n);
    for (int i = 0; i < n; i++)
        cin >> c[i], c[i]--, count[c[i]]++;
    vector<bool> vis(n);
    vector<int> sub(n);
    function<void(int, int)> dfs_size = [&](int src, int par) {
        sub[src] = 1;
        for (int nxt : adj[src]) {
            if (nxt == par || vis[nxt]) 
                continue;
            dfs_size(nxt, src);
            sub[src] += sub[nxt];
        }
    };
    auto get_centroid = [&](int src) -> int { 
        dfs_size(src, -1);
        int num = sub[src];
        int par = -1;
        do {
            int go = -1;
            for (int nxt : adj[src]) {
                if (nxt == par || vis[nxt]) 
                    continue;
                if (sub[nxt] * 2 > num)
                    go = nxt;
            }
            par = src;
            src = go;
        } while (src != -1);
        return par;
    };
    const int INF = 1e9;
    int ans = INF;
    vector<vector<int>> colors(n);
    vector<int> parent(n);
    vector<bool> done(n);
    vector<int> used;
    list<int> que;
    function<void(int)> solve = [&](int vert) {
        int root = get_centroid(vert);
        used.clear();
        function<void(int, int)> dfs_reset = [&](int src, int par) {
            parent[src] = -1;
            done[c[src]] = false;
            colors[c[src]].clear();
            for (int nxt : adj[src]) {
                if (nxt == par || vis[nxt])
                    continue;
                dfs_reset(nxt, src);
            }
        };
        function<void(int, int)> dfs_precompute = [&](int src, int par) {
            parent[src] = par;
            colors[c[src]].push_back(src);
            for (int nxt : adj[src]) {
                if (nxt == par || vis[nxt]) 
                    continue;
                dfs_precompute(nxt, src);
            }
        };
        dfs_reset(root, -1);
        dfs_precompute(root, -1);
        int tot = 0;
        auto add_color = [&](int col) {
            if (done[col])
                return;
            tot++;
            done[col] = true;
            used.push_back(col);
            for (int v : colors[col]) {
                que.push_back(v);
            }
        };
        add_color(c[root]);
        while (!que.empty()) {
            int top = que.front();
            que.pop_front();
            int par = parent[top];
            if (par != -1)
                add_color(c[par]);
        }
        bool ok = true;
        for (int x : used) 
            if ((int) colors[x].size() != count[x]) 
                ok = false;
        if (ok)
            ans = min(ans, tot - 1);
        vis[root] = true;
        for (int nxt : adj[root]) {
            if (vis[nxt]) 
                continue;
            solve(nxt);
        }
    };
    solve(0);
    cout << ans << '\n';
    return 0;
}  

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