LOJ3280「JOISC 2020 Day4」首都城市

题意

有一棵 \(n\) 个节点的数，每个节点有一个颜色（共 \(k\) 种），要求选出几种颜色，使这几种颜色的点构成联通块，求最小颜色数。

$ 1 \leq n,k \leq 2 \times 10^5$

思路

点分治，考虑经过分治重心的答案。

因为联通，所有点向重心的路径上的颜色都必须染，再将这些颜色加入，直到没有需要扩展的。注意一旦跳到一个已被染色的点即可跳出，因为上面的祖先已被考虑。注意如果出现某些点不属于当时分治区域，则不能记录到答案中。

#include <bits/stdc++.h>
using std::vector;
using std::queue;
const int N=200005;
int f[N],b[N],siz[N],vis[N],cvis[N],col[N],g,ans,now,n,m,x,y;
vector<int> e[N],c[N];
queue<int> q;
void dfs(int x,int fa){
	f[x]=fa,siz[x]=1;
	b[x]=now;
	for (auto u:e[x]){
		if (vis[u] || u==f[x]) continue;
		dfs(u,x);
		siz[x]+=siz[u];
	}
}
void find(int x,int size){
	int ret=size-siz[x];
	for (auto u:e[x]){
		if (vis[u] || u==f[x]) continue;
		find(u,size);
		ret=std::max(ret,siz[u]);
	}
	if (ret<=size/2) g=x;
}
void solve(int x){
	now++;
	dfs(x,0);
	cvis[col[x]]=now;
	int ff=1,cnt=0;
	while (!q.empty()) q.pop();
	q.push(col[x]);
	while (ff && !q.empty()){
		cnt++;
		int cx=q.front();
		q.pop();
		for (auto i:c[cx]){
			if (b[i]!=now){
				ff=0;
				break;
			}
			for (int j=f[i];j;j=f[j])
				if (cvis[col[j]]!=now){
					q.push(col[j]);
					cvis[col[j]]=now;
				}else break;
		}	
	}
	if (ff) ans=std::min(ans,cnt);	
	vis[x]=1;
	for (auto u:e[x]){
		if (vis[u]) continue;
		find(u,siz[u]);
		solve(g);
	}
}
int main(){
	scanf("%d%d",&n,&m);
	for (int i=1;i<n;i++){
		scanf("%d%d",&x,&y);
		e[x].push_back(y);
		e[y].push_back(x);
	}
	for (int i=1;i<=n;i++)
		scanf("%d",&col[i]),c[col[i]].push_back(i);
	ans=m;
	dfs(1,0);find(1,n);
	solve(g);
	printf("%d\n",ans-1);
}