Submission #839881

#	Time	Username	Problem	Language	Result	Execution time	Memory
839881		NK_	Capital City (JOI20_capital_city)	C++17	100 / 100	606 ms	45632 KiB

capital_city

// Success consists of going from failure to failure without loss of enthusiasm
#include <bits/stdc++.h>
 
using namespace std;
 
#define nl '\n'
#define pb push_back
#define pf push_front
 
#define mp make_pair
#define f first
#define s second
#define sz(x) int(x.size())
 
template<class T> using V = vector<T>;
using pi = pair<int, int>;
using vi = V<int>;
using vpi = V<pi>;
 
using ll = long long;
using pl = pair<ll, ll>;
using vpl = V<pl>;
using vl = V<ll>;
 
using db = long double;
 
template<class T> using pq = priority_queue<T, V<T>, less<T>>;
 
const int MOD = 1e9 + 7;
const ll INFL = ll(1e17);
const int LG = 19;

int main() {
	cin.tie(0)->sync_with_stdio(0);

	int N, K; cin >> N >> K;

	V<vi> adj(N), pos(K);

	for(int i = 0; i < N - 1; i++) {
		int u, v; cin >> u >> v; --u, --v;
		adj[u].pb(v); adj[v].pb(u);
	}

	vi A(N); for(auto& x : A) { cin >> x; --x; }

	vi proc(N), sub(N); 
	vi par(N), rt(N), vis(N); 
	vi seen(K), C(K), U;

	for(auto& x : A) C[x]++;

	function<void(int, int)> gen = [&](int u, int p) {
		sub[u] = 1;

		for(auto& v : adj[u]) if (!proc[v] && v != p) {
			gen(v, u);
			sub[u] += sub[v];
		}
	};

	function<void(int, int)> fix = [&](int u, int p) {
		par[u] = p; rt[u] = rt[p]; U.pb(u);
		for(auto& v : adj[u]) if (!proc[v] && v != p) fix(v, u);
	};

	function<int(int, int, int)> find = [&](int u, int p, int n) {
		for(auto& v : adj[u]) if (!proc[v] && v != p) {
			if (2 * sub[v] > n) return find(v, u, n);
		}
		return u;
	};


	int ans = MOD;
	function<void(int)> decomp = [&](int u) {
		gen(u, u); int c = find(u, u, sub[u]);
		rt[c] = c; fix(c, c);

		// cout << "NODES: ";
		for(auto& x : U) {
			// cout << x << " ";
			pos[A[x]].pb(x);
		}
		// cout << endl;


		// cout << "CENTROID: " << c << endl;

		vi have; 
		vis[c] = 1; have.pb(c); 

		int res = 0;
		for(int i = 0; i < sz(have); i++) {
			int x = have[i];
			// cout << "AT: " << x << endl;

			if (!seen[A[x]]) {
				++res;

				if (sz(pos[A[x]]) != C[A[x]]) {
					res = MOD; break;
				}

				for(auto& y : pos[A[x]]) if (!vis[y]) { 
					have.pb(y); vis[y] = 1;
				}
				
				seen[A[x]] = 1;
			}

			if (!vis[par[x]]) {
				have.pb(par[x]);
				vis[par[x]] = 1;
			}
		}

		for(auto& x : U) {
			vis[x] = 0, seen[A[x]] = 0;
			pos[A[x]].pop_back();
			// cout << A[x] << " -> " << sz(pos[A[x]]) << endl;
		}

		U = {};

		// cout << "RES: " << res << endl;

		ans = min(ans, res);

		proc[c] = 1;
		for(auto& v : adj[c]) if (!proc[v]) decomp(v);
	};

	decomp(0);

	cout << ans - 1 << nl;

	exit(0-0);
} 			

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