Submission #559475

#	Time	Username	Problem	Language	Result	Execution time	Memory
559475		dooompy	Capital City (JOI20_capital_city)	C++17	100 / 100	1159 ms	49252 KiB

capital_city

#include "bits/stdc++.h"

using namespace std;

void abc() {cout << endl;}
template <typename T, typename ...U> void abc(T a, U ...b) {
    cout << a << ' ', abc(b...);
}
template <typename T> void printv(T l, T r) {
    while (l != r) cout << *l << " \n"[++l == r];
}
template <typename A, typename B> istream& operator >> (istream& o, pair<A, B> &a) {
    return o >> a.first >> a.second;
}
template <typename A, typename B> ostream& operator << (ostream& o, pair<A, B> a) {
    return o << '(' << a.first << ", " << a.second << ')';
}
template <typename T> ostream& operator << (ostream& o, vector<T> a) {
    bool is = false;
    for (T i : a) {o << (is ? ' ' : '{'), is = true, o << i;}
    return o << '}';
}

#ifdef local
#define test(args...) abc("[" + string(#args) + "]", args)
#else
#define test(args...) void(0)
#endif

using ll = long long;

vector<int> adj[200005];
int city[200005];
vector<int> col[200005];
bool centroid[200005];
int sz[200005], par[200005];
vector<int> curcol[200005];
int ans; int n; int timer;
int seen[200005];

void dfs(int node, int p = -1) {

    sz[node] = 1;
    curcol[city[node]].push_back(node);

    for (auto nxt : adj[node]) {
        if (nxt == p || centroid[nxt]) continue;
        dfs(nxt, node);
        sz[node] += sz[nxt];
        par[nxt] = node;
    }
}

int find_centroid(int node, int p, int n) {

    for (auto nxt : adj[node]) {
        if (nxt == p || centroid[nxt]) continue;
        if (sz[nxt] * 2 >= n) return find_centroid(nxt, node, n);
    }

    return node;
}

void clear(int node, int p = -1) {
    curcol[city[node]].clear();

    for (auto nxt : adj[node]) {
        if (nxt == p || centroid[nxt]) continue;
        clear(nxt, node);
    }
}

void solve(int node) {
    dfs(node);

    int cent = find_centroid(node, -1, sz[node]);

    clear(node);

    dfs(cent);

    queue<int> q;
    q.push(city[cent]);

    int ct = 0;
    timer++;

    seen[city[cent]] = timer;

    while (!q.empty()) {
        ct++;
        auto cur = q.front(); q.pop();
        if (curcol[cur].size() != col[cur].size()) {
            ct = n; break;
        }

        for (auto nd : curcol[cur]) {
            if (nd == cent) continue;

            int p = city[par[nd]];
            if (seen[p] == timer) continue;
            seen[p] = timer;
            q.push(p);
        }
    }

    ans = min(ans, ct - 1);

    clear(cent);
    centroid[cent] = true;

    for (auto nxt : adj[cent]) {
        if (!centroid[nxt]) solve(nxt);
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
//    freopen("", "r", stdin);
//    freopen("", "w", stdout);
    int k; cin >> n >> k; ans = n;

    for (int i = 0; i < n-1; i++) {
        int a, b; cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    for (int i = 1; i <= n; i++) {
        cin >> city[i];
        col[city[i]].push_back(i);
    }

    solve(1);

    cout << ans << "\n";
}

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