Submission #216338

#	Time	Username	Problem	Language	Result	Execution time	Memory
216338		abacaba	Capital City (JOI20_capital_city)	C++14	100 / 100	939 ms	35068 KiB

capital_city

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

#define pb push_back
#define ll long long
#define pii pair<int, int>

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

template <typename T> inline T range(T l, T r) {
	return uniform_int_distribution <T>(l, r)(rng);
}

const int N = 2e5 + 15;
const int inf = 2e9;
const int mod = 1e9 + 7;
int n, m, k, a[N];
vector <int> g[N];
bool used[N];
int total[N], cnt[N];
int cl, sz[N], now, ans = inf;

int curp[N];
int q[N], ql, qr;

vector <int> ord[N];

struct dsu {
	int par[N];
	int find(int v) {
		if(v == par[v])
			return v;
		return par[v] = find(par[v]);
	}
	bool unio(int a, int b) {
		a = find(a);
		b = find(b);
		if(a == b)
			return false;
		if(range(0, 1))
			swap(a, b);
		par[b] = a;
		return true;
	}
} st;

void dfs_sz(int v, int p) {
	sz[v] = 1;
	st.par[a[v]] = a[v];
	ord[a[v]].clear();
	for(int to : g[v])
		if(to != p && !used[to]) {
			curp[to] = v;
			dfs_sz(to, v);
			sz[v] += sz[to];
		}
}

int centroid(int v, int p, int n) {
	for(int to : g[v])
		if(to != p && !used[to] && 2 * sz[to] > n)
			return centroid(to, v, n);
	return v;
}

void calc(int v, int p) {
	ord[a[v]].pb(v);
	for(int to : g[v])
		if(to != p && !used[to])
			calc(to, v);
}

void solve(int c) {
	curp[c] = -1;
	used[c] = true;
	dfs_sz(c, c);
	calc(c, c);

	ql = 0, qr = -1;

	for(int x : ord[a[c]])
		q[++qr] = x;

	int cur = 0, now = 0;

	while(ql <= qr) {
		int v = q[ql++];
		cur += (int)ord[a[v]].size() - total[a[v]];
		// if(c == 2)
		// 	cout << ord[a[v]].size() << ' ' << total[a[v]] << endl;
		if(curp[v] == -1)
			continue;
		if(st.unio(a[curp[v]], a[v])) {
			++now;
			for(int x : ord[a[curp[v]]])
				q[++qr] = x;
		}
	}
	if(cur == 0 && now < ans)
		ans = now;
	for(int to : g[c])
		if(!used[to])
			solve(centroid(to, c, sz[to]));
}

main() {
	ios_base::sync_with_stdio(0); cout.tie(0); cin.tie(0);
	cin >> n >> k;
	for(int i = 1; i < n; ++i){
		int u, v;
		cin >> u >> v;
		g[u].pb(v);
		g[v].pb(u);
	}
	for(int i = 1; i <= n; ++i) {
		cin >> a[i];
		++total[a[i]];
	}
	dfs_sz(1, 1);
	solve(centroid(1, 1, n));
	cout << ans << endl;
	return 0;
}

Compilation message (stderr)

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

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