Submission #216922

#	Time	Username	Problem	Language	Result	Execution time	Memory
216922		JustasZ	Capital City (JOI20_capital_city)	C++14	100 / 100	1162 ms	34520 KiB

capital_city

/*input
*/
#include <bits/stdc++.h>
#define pb push_back
#define x first
#define y second
#define all(a) a.begin(), a.end()
#define sz(a) (int)a.size()
#define rd() abs((int)rng())
using namespace std;
using ll = long long;
using ld = long double;
using pii = pair<int, int>;
const int maxn = 2e5 + 100;
const int mod = 1e9 + 7;
mt19937 rng(chrono::high_resolution_clock::now().time_since_epoch().count());
int n, k, color[maxn], vis[maxn], par[maxn], par2[maxn], sub[maxn], color_added[maxn], color_id[maxn], color_cnt[maxn], color_total[maxn];
vector<int>byColor[maxn];
vector<int>colors;
vector<int>adj[maxn];
void dfs(int v) {
	sub[v] = 1;
	for (int to : adj[v]) {
		if (!vis[to] && to != par[v]) {
			par[to] = v;
			dfs(to);
			sub[v] += sub[to];
		}
	}
}
int get_centroid(int v) {
	par[v] = 0;
	dfs(v);
	int siz = sub[v];
	while (true) {
		pii mx = {0, 0};
		for (int to : adj[v]) {
			if (!vis[to] && to != par[v]) {
				mx = max(mx, {sub[to], to});
			}
		}
		if (mx.x * 2 > siz) {
			v = mx.y;
		} else {
			return v;
		}
	}
}
void dfs2(int v, int p) {
	par2[v] = p;
	if (color_id[color[v]] == -1) {
		color_id[color[v]] = sz(colors);
		colors.pb(color[v]);
	}
	for (int to : adj[v]) {
		if (!vis[to] && to != p) {
			dfs2(to, v);
		}
	}
}
void dfs3(int v, int p) {
	byColor[color_id[color[v]]].pb(v);
	++color_cnt[color[v]];
	for (int to : adj[v]) {
		if (!vis[to] && to != p) {
			dfs3(to, v);
		}
	}
}
int ans = mod;
void solve(int v, int level = 0) {
	par[v] = 0;
	v = get_centroid(v);
	vis[v] = 1;
	// cout << "level: " << level << ", vertex: " << v << "\n";

	for (int c : colors) {
		if (color_id[c] != -1) {
			byColor[color_id[c]].clear();
		}
		color_cnt[c] = 0;
		color_id[c] = -1;
		color_added[c] = 0;
	}
	colors.clear();

	dfs2(v, 0);
	dfs3(v, 0);

	queue<int>Q;
	Q.push(color[v]);
	color_added[color[v]] = 1;

	int total_colors = 1;
	while (sz(Q)) {
		int c = Q.front();
		if (color_cnt[c] != color_total[c]) {
			total_colors = mod;
		}
		Q.pop();
		for (int u : byColor[color_id[c]]) {
			int p = par2[u];
			if (p != 0 && !color_added[color[p]]) {
				++total_colors;
				color_added[color[p]] = 1;
				Q.push(color[p]);
			}
		}
	}

	ans = min(ans, total_colors - 1);

	for (int to : adj[v]) {
		if (!vis[to]) {
			solve(to, level + 1);
		}
	}
}
int fast() {
	for (int i = 1; i <= k; i++) {
		color_total[i] = 0;
		color_id[i] = -1;
	}
	for (int i = 1; i <= n; i++) {
		++color_total[color[i]];
		vis[i] = 0;
	}
	ans = mod;
	solve(1);
	return ans;
}
int main()
{
	ios_base::sync_with_stdio(false), cin.tie(0);
	cin >> n >> k;
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		cin >> a >> b;
		adj[a].pb(b);
		adj[b].pb(a);
	}

	for (int i = 1; i <= n; i++) {
		cin >> color[i];
	}

	cout << fast() << "\n";
	return 0;
}

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