Submission #367452

#	Time	Username	Problem	Language	Result	Execution time	Memory
367452		valerikk	Capital City (JOI20_capital_city)	C++17	100 / 100	1063 ms	39520 KiB

capital_city

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

#define pb push_back

const int MXN = 2e5 + 7;

int n, k;
vector<int> ed[MXN];
int c[MXN], tot[MXN], sz[MXN], ans = MXN;
bool del[MXN], used[MXN];
vector<int> nodes;
bool usedc[MXN];
vector<int> cc[MXN];
int pp[MXN];

int get_sz(int v, int p = -1) {
    sz[v] = 1;
    for (int u : ed[v]) {
        if (u != p && !del[u]) sz[v] += get_sz(u, v);
    }
    return sz[v];
}

int cent(int v, int nn, int p = -1) {
    for (int u : ed[v]) {
        if (u != p && !del[u] && 2 * sz[u] > nn) return cent(u, nn, v);
    }
    return v;
}

void dfs(int v, int p = -1) {
    pp[v] = p;
    nodes.pb(v);
    used[v] = usedc[c[v]] = 0;
    cc[c[v]].pb(v);
    for (int u : ed[v]) {
        if (u != p && !del[u]) dfs(u, v);
    }
}

void solve(int v) {
    get_sz(v);
    nodes.clear();
    dfs(v);
    queue<int> q;
    q.push(c[v]);
    used[v] = usedc[c[v]] = 1;
    int need = 0;
    while (!q.empty()) {
        int x = q.front();
        q.pop();
        for (int u : cc[x]) {
            int vv = u;
            while (!used[vv]) {
                used[vv] = 1;
                if (!usedc[c[vv]]) {
                    need++;
                    usedc[c[vv]] = 1;
                    q.push(c[vv]);
                }
                vv = pp[vv];
            }
        }
    }
    bool ok = 1;
    for (int u : nodes) {
        if (used[u]) ok &= cc[c[u]].size() == tot[c[u]];
    }
    for (int u : nodes) cc[c[u]].clear();
    if (ok) ans = min(ans, need);
    del[v] = 1;
    for (int u : ed[v]) {
        if (!del[u]) solve(cent(u, sz[u]));
    }
}   



int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> k;
    for (int i = 0; i < n - 1; i++) {
        int a, b;
        cin >> a >> b;
        a--;
        b--;
        ed[a].pb(b);
        ed[b].pb(a);
    }
    for (int i = 0; i < n; i++) {
        cin >> c[i];
        c[i]++;
        tot[c[i]]++;
    }
    solve(cent(0, get_sz(0)));
    cout << ans << endl;
    return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'void solve(int)':
capital_city.cpp:68:44: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   68 |         if (used[u]) ok &= cc[c[u]].size() == tot[c[u]];
      |                            ~~~~~~~~~~~~~~~~^~~~~~~~~~~~

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