maki和tree （并查集）

maki和tree

#include<bits/stdc++.h>
using namespace std;
#define ll long long 
int n;
char color[111111];
int fa[111111];
int kdn[111111];//孩子数目
ll sum[111111];//白色连通块数目 
vector<int> G[111111];
int f(int x)
{
	if(fa[x]==x) return x;
	return f(fa[x]);
}
void uni(int x,int y)
{
	x=f(x) ,y=f(y);//!!! 父节点保存！！ 
	if(x==y) return ;
	if(kdn[x]<kdn[y])
	{
		fa[x]=y;
		kdn[y]+=kdn[x]+1;
		return ;
	}
	else 
	{
		fa[y]=x;
		kdn[x]+=kdn[y]+1;
		return ;
	}
}
ll countt(vector<int>T) //T[i]为 第i个相邻白书连通块数目 
{	
	ll res=0;//黑色节点在端点的情况 = 白色点数目 
 	ll res2=0;//黑色节点在中间的情况 
	int size=T.size();
	if(size==0 ) return 0;
	ll summ[size+10];
	summ[0]=0;
	
	for(int i=0;i<size;i++)
	{
		res+=T[i];
		summ[i]=res;
	}
	
	for(int i=1;i<size;i++)
	{
		res2+=summ[i-1]*T[i];
	}
	return res+res2;
	
}
int main()
{
	scanf("%d",&n);
	scanf("%s",color+1);
	for(int i=1;i<=n;i++) fa[i]=i;
	for(int i=0;i<n-1;i++)
	{
		int x,y;
		scanf("%d%d",&x,&y);
		G[x].push_back(y);
		G[y].push_back(x);
		if(color[x]=='W'&&color[y]=='W')
		{
			uni(x,y);
		}
	}
	for(int i=1;i<=n;i++) sum[i]=kdn[f(i)]+1;//!!!kdn[fa(i)]  加上父亲节点的孩子数 
	ll ans=0;
	for(int i=1;i<=n;i++) 
	{
		if(color[i]=='B')
		{
			int size=G[i].size();
			vector<int> T;
			T.clear();
			for(int j=0;j<size;j++)
			{
				if(color[G[i][j]]=='W') T.push_back(sum[G[i][j]]);//将连通块大小放入 				
			}
			ans+=countt(T);
		}
	}
	cout<<ans<<endl;	
	
}