【圆方树】【启发式合并】CodeChef Chef and Sad Pairs

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/qq_34454069/article/details/88809062

分析：

圆方树板子题

每个点维护一下它子树中的颜色。

启发式合并算贡献

不过也可以分颜色用虚树做

#include<cstdio>
#include<cstring>
#include<vector>
#include<map>
#include<algorithm>
#define SF scanf
#define PF printf
#define MAXN 400010
#define MAXG 1000010
using namespace std;
typedef long long ll;
typedef map<int,int>::iterator mit;
int n,m;
int dfn[MAXN],low[MAXN],cnt;
int st[MAXN],tp;
struct node{
	int u,v;
	node () {}
	node(int u1,int v1):u(u1),v(v1) {}	
}Edge[MAXN];
int tot,n2;
vector<int> a[MAXN];
void dfs(int x,int fa=0){
	dfn[x]=low[x]=++cnt;
	st[++tp]=x;
	for(int i=0;i<int(a[x].size());i++){
		int u=a[x][i];
		if(u==fa)
			continue;
		if(dfn[u]==0){
			dfs(u,x);
			low[x]=min(low[x],low[u]);
			if(low[u]>=dfn[x]){
				n2++;
				Edge[++tot]=node(n2,x);
				while(1){
					Edge[++tot]=node(n2,st[tp]);
					if(st[tp--]==u)
						break;
				}
			}
		}
		else
			low[x]=min(low[x],dfn[u]);
	}
}
int idx[MAXN],col[MAXN];
map<int,int> mp[MAXN];
int totc[MAXG],tot1[MAXG],son[MAXN];
ll sum[MAXN],ans[MAXN];
void add(int x,int c,int s){
	sum[x]-=1ll*mp[idx[x]][c]*(totc[c]-mp[idx[x]][c]);
	mp[idx[x]][c]+=s;
	ans[x]+=1ll*s*(totc[c]-mp[idx[x]][c]);
	sum[x]+=1ll*mp[idx[x]][c]*(totc[c]-mp[idx[x]][c]);
}
bool vis[MAXN];
void prepare(int x){
	vis[x]=1;
	if(col[x])
		totc[col[x]]++;
	for(int i=0;i<int(a[x].size());i++){
		int u=a[x][i];
		if(vis[u])
			continue;
		prepare(u);	
	}
}
void solve(int x,int fa=0){
	for(int i=0;i<int(a[x].size());i++){
		int u=a[x][i];
		if(u==fa)
			continue;
		solve(u,x);
		if(son[x]==0||mp[idx[u]].size()>mp[idx[son[x]]].size())
			son[x]=u;
	}
	if(son[x]==0)
		idx[x]=++tot;
	else{
		idx[x]=idx[son[x]];
		sum[x]=sum[son[x]];
		ans[x]+=sum[x];
	}
	if(x<=n)
		add(x,col[x],1);
//	PF("{%d %lld %lld}\n",x,sum[x],ans[x]);
	for(int i=0;i<int(a[x].size());i++){
		int u=a[x][i];
		if(u==fa||u==son[x])
			continue;
		for(map<int,int>::iterator it=mp[idx[u]].begin();it!=mp[idx[u]].end();it++)
			add(x,it->first,it->second);
	}
//	PF("[%d]-\n",x);
//	for(map<int,int>::iterator it=mp[idx[x]].begin();it!=mp[idx[x] ].end();it++)
//		PF("<%d,%d>\n",it->first,it->second);
}
int main(){
	SF("%d%d",&n,&m);
	for(int i=1;i<=n;i++)
		SF("%d",&col[i]);
	int u,v;
	for(int i=1;i<=m;i++){
		SF("%d%d",&u,&v);
		a[u].push_back(v);
		a[v].push_back(u);
	}
	n2=n;
	for(int i=1;i<=n;i++)
		if(dfn[i]==0)
			dfs(i);
	for(int i=1;i<=n;i++) a[i].clear();
	for(int i=1;i<=tot;i++){
		a[Edge[i].u].push_back(Edge[i].v);
		a[Edge[i].v].push_back(Edge[i].u);	
	}
	tot=0;
	ll ansx=0;
	for(int i=1;i<=n2;i++)
		if(vis[i]==0){
			prepare(i);
			solve(i);
			for(mit it=mp[idx[i]].begin();it!=mp[idx[i]].end();it++){
				ansx+=1ll*tot1[it->first]*it->second;
				tot1[it->first]+=it->second;
				totc[it->first]=0;
			}
		}
	for(int i=1;i<=n;i++)
		PF("%lld\n",ans[i]+ansx);
	
}