Submission #538557

#	Time	Username	Problem	Language	Result	Execution time	Memory
538557		pavement	Capital City (JOI20_capital_city)	C++17	100 / 100	863 ms	45580 KiB

capital_city

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#ifdef _WIN32
#define getchar_unlocked _getchar_nolock
#endif
#define int long long
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define ppb pop_back
#define eb emplace_back
#define g0(a) get<0>(a)
#define g1(a) get<1>(a)
#define g2(a) get<2>(a)
#define g3(a) get<3>(a)
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
using db = double;
using ll = long long;
using ld = long double;
using ii = pair<int, int>;
using iii = tuple<int, int, int>;
using iiii = tuple<int, int, int, int>;
template<class key, class value = null_type, class cmp = less<key> >
using ordered_set = tree<key, value, cmp, rb_tree_tag, tree_order_statistics_node_update>;

int N, K, ans = LLONG_MAX, root, tot_sz, sz[200005], C[200005], par[200005];
bool vis[200005], seen[200005], proc[200005], in_comp[200005];
vector<int> vec_tmp, adj[200005], vec[200005];
queue<int> Q;

int get_sz(int n, int e = -1) {
	sz[n] = 1;
	for (auto u : adj[n]) if (u != e && !proc[u])
		sz[n] += get_sz(u, n);
	return sz[n];
}

void get_centroid(int n, int e = -1) {
	int m = 0;
	for (auto u : adj[n]) if (u != e && !proc[u]) {
		get_centroid(u, n);
		m = max(m, sz[u]);
	}
	m = max(m, tot_sz - sz[n]);
	if (m <= tot_sz / 2) root = n;
}

void init(int n, int e = -1) {
	vec_tmp.pb(n);
	vis[n] = seen[C[n]] = 0;
	par[n] = e;
	for (auto u : adj[n]) if (u != e && !proc[u]) init(u, n);
}

void decomp(int n) {
	get_sz(n);
	tot_sz = sz[n];
	get_centroid(n);
	// root is the new centroid
	vec_tmp.clear();
	init(root);
	while (!Q.empty()) Q.pop();
	for (auto i : vec_tmp) in_comp[i] = 1;
	Q.push(C[root]);
	seen[C[root]] = 1;
	int mrg = 0;
	bool inv = 0;
	while (!Q.empty()) {
		int x = Q.front();
		mrg++;
		Q.pop();
		for (auto i : vec[x]) {
			if (!in_comp[i]) {
				inv = 1;
				goto done;
			}
			while (!vis[i]) {
				if (!seen[C[i]]) {
					seen[C[i]] = 1;
					Q.push(C[i]);
				}
				vis[i] = 1;
				if (par[i] == -1) break;
				i = par[i];
			}
		}
	}
	done:;
	for (auto i : vec_tmp) in_comp[i] = 0;
	if (!inv) ans = min(ans, mrg - 1);
	proc[root] = 1;
	for (auto u : adj[root])
		if (!proc[u]) decomp(u);
}

main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> N >> K;
	for (int i = 1, u, v; i < N; i++) {
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	for (int i = 1; i <= N; i++) {
		cin >> C[i];
		vec[C[i]].pb(i);
	}
	decomp(1);
	cout << ans << '\n';
}

Compilation message (stderr)

capital_city.cpp:99:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   99 | main() {
      | ^~~~

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