Submission #985275

#	Time	Username	Problem	Language	Result	Execution time	Memory
985275		alextodoran	Capital City (JOI20_capital_city)	C++17	100 / 100	454 ms	42444 KiB

capital_city

/**
 _  _   __  _ _ _  _  _ _
 |a  ||t  ||o    d | |o  |
| __    _| | _ | __|  _ |
| __ |/_  | __  /__\ / _\|

**/

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int N_MAX = 200000;

int N, K;
vector <int> adj[N_MAX + 2];
int C[N_MAX + 2];

bool erased[N_MAX + 2];

vector <int> nodes;
int parent[N_MAX + 2];
int sub[N_MAX + 2];
bool in_comp[N_MAX + 2];

vector <int> with_color[N_MAX + 2];

void dfs (int u, bool first) {
    if (first == true) {
        nodes.push_back(u);
        in_comp[u] = true;
    }
    sub[u] = 1;
    for (int v : adj[u]) {
        if (v != parent[u] && erased[v] == false) {
            parent[v] = u;
            dfs(v, first);
            sub[u] += sub[v];
        }
    }
}

int answer = INT_MAX / 2;
bool seen[N_MAX + 2];
bool cleared[N_MAX + 2];

void solve (int root) {
    parent[root] = 0; dfs(root, true);
    int centroid = 0;
    for (int u : nodes) {
        if (sub[root] - sub[u] > (int) nodes.size() / 2) {
            continue;
        }
        bool ok = true;
        for (int v : adj[u]) {
            if (v != parent[u] && erased[v] == false && sub[v] > (int) nodes.size() / 2) {
                ok = false;
                break;
            }
        }
        if (ok == true) {
            centroid = u;
            break;
        }
    }
    root = centroid;
    parent[root] = 0; dfs(root, false);
    queue <int> q;
    q.push(C[root]);
    cleared[C[root]] = true;
    int merges = 0;
    while (q.empty() == false) {
        int c = q.front(); q.pop();
        for (int u : with_color[c]) {
            if (in_comp[u] == false) {
                merges = INT_MAX / 2;
                break;
            }
            while (u != 0 && seen[u] == false) {
                if (cleared[C[u]] == false) {
                    cleared[C[u]] = true;
                    q.push(C[u]);
                    merges++;
                }
                seen[u] = true;
                u = parent[u];
            }
        }
    }
    answer = min(answer, merges);
    for (int u : nodes) {
        in_comp[u] = false;
        seen[u] = false;
        cleared[C[u]] = false;
    }
    nodes.clear();
    erased[root] = true;
    for (int u : adj[root]) {
        if (erased[u] == false) {
            solve(u);
        }
    }
}

int main () {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    cin >> N >> K;
    for (int i = 1; i <= N - 1; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for (int u = 1; u <= N; u++) {
        cin >> C[u];
        with_color[C[u]].push_back(u);
    }
    solve(1);
    cout << answer << "\n";

    return 0;
}

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