Submission #945174

#	Time	Username	Problem	Language	Result	Execution time	Memory
945174		ArashJafariyan	Capital City (JOI20_capital_city)	C++17	100 / 100	473 ms	42072 KiB

capital_city

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

#ifdef DEBUG
#include "debug.h"
#else
#define debug(...) 0
#endif

#define pb push_back
#define i2 array<int, 2>
#define ll long long

const int N = 2e5 + 4;

int n, k, c[N], sz[N], root[N], par[N];
bool done[N], in[N];
vector<int> adj[N], V[N];
queue<int> q;

void fill_sz(int v, int p = 0) {
    sz[v] = 1;
    for (int u : adj[v]) {
        if (u != p && !done[u]) {
            fill_sz(u, v);
            sz[v] += sz[u];
        }
    }
}

int find_cen(int v, int siz, int p = 0) {
    for (int u : adj[v]) {
        if (!done[u] && u != p && sz[u] > siz / 2) {
            return find_cen(u, siz, v);
        }
    }
    return v;
}

void fill(int v, int p = 0) {
    in[c[v]] = 0;
    par[v] = p;
    root[v] = (p == 0 ? v : root[p]);
    for (int u : adj[v]) {
        if (u != p && !done[u]) {
            fill(u, v);
        }
    }
}

int dfs(int v = 1) {
    fill_sz(v);
    v = find_cen(v, sz[v]);
    fill(v);
    done[v] = true;
    q.push(c[v]);

    int ans = 0;
    while (q.size()) {
        int col = q.front();
        q.pop();
        
        if (in[col] == 1) {
            continue;
        }
        in[col] = 1;
        for (int u : V[col]) {
            if (root[u] != v) {
                ans = k;
                break;
            }
            if (par[u] != 0) {
                q.push(c[par[u]]);
            }    
        }
        ans++;
    }

    for (int u : adj[v]) {
        if (!done[u]) {
            ans = min(ans, dfs(u));
        }
    }
    return ans;
}

void solve() {
    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int v, u;
        cin >> v >> u;
        adj[v].pb(u);
        adj[u].pb(v);
    }
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
        V[c[i]].pb(i);
    }
    cout << dfs() - 1;
}

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

    int t = 1;
    // cin >> t;
    while (t--) {
        solve();
    }

    return 0;
}

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