Submission #219299

#	Time	Username	Problem	Language	Result	Execution time	Memory
219299		maruii	Capital City (JOI20_capital_city)	C++14	100 / 100	891 ms	33132 KiB

capital_city

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

int N, K, C[200001], ans, A[200001], B[200001];
vector<int> edge[200001];

bool vis[200001];
int sz[200001];
int szf(int x, int p) {
	sz[x] = 1;
	for (auto i : edge[x]) if (i != p && !vis[i]) sz[x] += szf(i, x);
	return sz[x];
}
int get_cen(int x, int p, int t) {
	for (auto i : edge[x]) if (i != p && !vis[i] && sz[i] >= t) return get_cen(i, x, t);
	return x;
}

bool ban[200001], chk[200001], vis2[200001];
vector<int> clr, vec[200001];
int par[200001];
void dfs(int x, int p) {
	par[x] = p;
	clr.eack(C[x]);
	vec[C[x]].eack(x);
	++B[C[x]];
	chk[C[x]] = vis2[x] = 0;
	for (auto i : edge[x]) if (i != p && !vis[i]) dfs(i, x);
}

void cd(int x) {
	x = get_cen(x, x, szf(x, x) + 1 >> 1);
	vis[x] = 1;
	for (auto i : clr) {
		vec[i].clear();
		B[i] = ban[i] = 0;
	}
	clr.clear();
	dfs(x, x);
	for (auto i : clr) if (A[i] > B[i]) ban[i] = 1;
	queue<int> q;
	q.push(C[x]);
	chk[C[x]] = vis2[x] = 1;
	bool flag = 1;
	int cnt = 0;
	while (q.size()) {
		int x = q.front(); q.pop();
		if (ban[x]) {
			flag = 0;
			break;
		}
		for (auto i : vec[x]) {
			for (int j = i; !vis2[j]; j = par[j]) {
				vis2[j] = 1;
				if (!chk[C[j]]) {
					q.push(C[j]);
					chk[C[j]] = 1;
					++cnt;
				}
			}
		}
	}
	if (flag) ans = min(ans, cnt);
	for (auto i : edge[x]) if (!vis[i]) cd(i);
}

int main() {
	ios_base::sync_with_stdio(0), cin.tie(0);
	cin >> N >> K;
	ans = K - 1;
	for (int i = 1; i < N; ++i) {
		int x, y; cin >> x >> y;
		edge[x].eack(y);
		edge[y].eack(x);
	}
	for (int i = 1; i <= N; ++i) {
		cin >> C[i];
		++A[C[i]];
	}
	cd(1);
	cout << ans;
	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'void cd(int)':
capital_city.cpp:33:30: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
  x = get_cen(x, x, szf(x, x) + 1 >> 1);
                    ~~~~~~~~~~^~~

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