Submission #1031111

#	Time	Username	Problem	Language	Result	Execution time	Memory
1031111		Blagoj	Capital City (JOI20_capital_city)	C++17	100 / 100	405 ms	42004 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

#define endl '\n'
#define ll long long
#define all(x) (x).begin(), (x).end()

const int mxn = 2e5 + 100;

int n, k;

vector<int> g[mxn], type(mxn), sz(mxn), p(mxn), cnt(mxn), have[mxn];
bool done[mxn], merged[mxn];

int getSize(int cur, int par = -1) {
    sz[cur] = 1;
    for (auto x : g[cur]) {
        if (x != par && !done[x]) sz[cur] += getSize(x, cur);
    }
    return sz[cur];
}

int getCentroid(int cur, int des, int par = -1) {
    for (auto x : g[cur]) {
        if (sz[x] >= des && x != par && !done[x]) return getCentroid(x, des, cur);
    }
    return cur;
}

void Clear(int cur, int par) {
    p[cur] = par;
    cnt[type[cur]] = merged[type[cur]] = 0;
    for (auto x : g[cur]) {
        if (x != par && !done[x]) Clear(x, cur);
    }
}

void get(int cur, int par) {
    cnt[type[cur]]++;
    for (auto x : g[cur]) {
        if (x != par && !done[x]) get(x, cur);
    }
}

ll ans = 1e12;

void decompose(int cur = 1) {
    cur = getCentroid(cur, getSize(cur) / 2);
    done[cur] = 1;
    Clear(cur, cur);
    get(cur, cur);
    stack<int> s;
    s.push(type[cur]);
    ll curAns = 0;
    while (s.size()) {
        int top = s.top();
        s.pop();
        if (merged[top]) continue;
        if (cnt[top] != have[top].size()) {
            curAns = 1e12;
            break;
        }
        curAns++;
        merged[top] = 1;
        for (auto x : have[top]) s.push(type[p[x]]);
    }
    ans = min(ans, curAns - 1);
    for (auto x : g[cur]) if (!done[x]) decompose(x);
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cin >> n >> k;
    for (int i = 0; i < n - 1; i++) {
        int f, t;
        cin >> f >> t;
        g[f].push_back(t);
        g[t].push_back(f);
    }
    for (int i = 1; i <= n; i++) {
        cin >> type[i];
        have[type[i]].push_back(i);
    }
    decompose();
    cout << ans;
}

Compilation message (stderr)

capital_city.cpp: In function 'void decompose(int)':
capital_city.cpp:60:22: warning: comparison of integer expressions of different signedness: '__gnu_cxx::__alloc_traits<std::allocator<int>, int>::value_type' {aka 'int'} and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   60 |         if (cnt[top] != have[top].size()) {

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