P10180 题解

P10180

思路

首先答案是所有连通块大小的平方和。

先用并查集将相同颜色合起来，记录 \(sum_i\) 表示只看颜色 \(i\) 的答案。

如果对于询问 \(u,v\) 不存在 \(e(i,j)\) 满足 \(a_i=u,a_j=v\)，答案为两个颜色单独的答案之和。

否则只有至多 \(n-1\) 对 \((u,v)\) 的至多 \(n-1\) 条边，需要考虑两种颜色合起来。可以用可撤销并查集，枚举所有对应的颜色对，将对应的边加入并查集，记录答案，再删去。复杂度 \(O(n\log n)\)。只要不是写的太臭就可以。

code

int n,q,a[maxn];
int f[maxn],siz[maxn];
int fd(int x){
	if(x==f[x])return x;
	return fd(f[x]);
}
int st[maxn][2],tp;
void merge(int u,int v){
	if(u==v)return ;
	if(siz[u]<siz[v])swap(u,v);
	st[++tp][0]=v,st[tp][1]=siz[v];
	f[v]=u;siz[u]+=siz[v];
}
#define pii pair<int,int>
unordered_map<int,int>mp;pii id[maxn];
vector<pii> ans[maxn];int idx;
int res[maxn],sum[maxn];
int get(int u,int v){return min(u,v)*n+max(u,v);}
void work(){
	n=read();q=read();
	for(int i=1;i<=n;i++)a[i]=read();
	for(int i=1;i<=n;i++)f[i]=i,siz[i]=1;
	for(int i=2;i<=n;i++){
		int u=read();
		if(a[u]==a[i])merge(fd(u),fd(i));
		else{
			int p=get(a[u],a[i]);
			if(!mp[p])mp[p]=++idx,id[idx]={a[u],a[i]};
			ans[mp[p]].push_back(make_pair(u,i));
		}
	}
	for(int i=1;i<=n;i++)if(f[i]==i)sum[a[i]]+=siz[i]*siz[i];
	for(int i=1;i<=idx;i++){
		// cout<<id[i].first<<" "<<id[i].second<<" a\n";
		res[i]=sum[id[i].first]+sum[id[i].second];
		int lst=tp;
		for(pii j:ans[i]){
			int u=fd(j.first),v=fd(j.second);
			res[i]+=2*siz[u]*siz[v];
			merge(u,v);
		}
		while(tp!=lst){
			siz[f[st[tp][0]]]-=st[tp][1];
			f[st[tp][0]]=st[tp][0];
			tp--;
		}
		// for(int j=1;j<=n;j++)if(j==f[j])cout<<j<<" "<<f[j]<<" "<<siz[j]<<"\n";
	}
	while(q--){
		int u=read(),v=read();
		if(mp.count(get(u,v)))printf("%lld\n",res[mp[get(u,v)]]);
		else printf("%lld\n",sum[u]+sum[v]);
	}
}