[ds记录]「JOISC 2020 Day4」首都城市

\(\texttt{link}\)

考虑点分治：每次计算以分治中心的城市为首都的合并次数，当已合并的城市的某些小镇不在当前这一层分治里，则必然会经过上一层的分治中心，答案也不会更优，直接忽略这一层。

具体实现：

首先遍历这一部分得到每个点的父亲节点，然后将分治中心所在城市的所有小镇加入一个队列里，如果有小镇不在这一层分治里，直接返回；否则的话将队列里的每个点的父亲加入队列，重复这个操作即可。

时间复杂度 \(O(n \log n)\)。

\(\texttt{Code:}\)

#include <bits/stdc++.h>
#define pb push_back
using namespace std;
const int N = 2e5 + 5;
int n, k, sum, ans, tot, rt, s[N], mn, vis[N], t[N], col[N], tag[N], fl, f[N];
vector<int> g[N], id[N];
queue<int> q;
void frt(int u, int fa) {
	s[u] = 1;
	int cur = 1;
	for (int v : g[u]) {
		if (v == fa || vis[v]) continue;
		frt(v, u);
		s[u] += s[v];
		cur = max(cur, s[v]);
	}
	cur = max(cur, tot - s[u]);
	if (cur < mn) mn = cur, rt = u;
}
void modify(int u, int fa, int x) {
	tag[u] = x; f[u] = fa;
	t[col[u]] = 0; s[u] = 1;
	for (int v : g[u]) {
		if (v == fa || vis[v]) continue;
		modify(v, u, x);
		s[u] += s[v];
	}
}
void insert(int u) {
	t[u] = 1; ++sum;
	for (int v : id[u]) {
		if (!tag[v]) {fl = 1; return;}
		q.push(v);
	}
}
void solve(int u) {
	vis[u] = 1;
	modify(u, 0, 1);
	while (!q.empty()) q.pop();
	fl = 0; sum = 0;
	insert(col[u]);
	while (!fl && !q.empty()) {
		int cur = q.front(); q.pop();
		if (cur == u) continue;
		if (!t[col[f[cur]]]) insert(col[f[cur]]);
	}
	if (!fl) ans = min(ans, sum - 1);
	modify(u, 0, 0);
	for (int v : g[u]) {
		if (vis[v]) continue;
		tot = s[v]; mn = tot + 1;
		frt(v, u); solve(rt);
	}
}
int main() {
	scanf("%d%d", &n, &k);
	ans = k - 1;
	for (int i = 1; i < n; i++) {
		int u, v; scanf("%d%d", &u, &v);
		g[u].pb(v); g[v].pb(u);
	}
	for (int i = 1; i <= n; i++) {
		scanf("%d", &col[i]);
		id[col[i]].pb(i);
	}
	tot = n; mn = tot + 1;
	frt(1, 0); solve(rt);
	printf("%d", ans);
    return 0;
}