WA on test 2

答案略大

#include<bits/stdc++.h>
#define int long long
#define MAXN 200005 
using namespace std;
int n,m,x,y,z,tmp[MAXN],p[MAXN],cnt,vis[MAXN],siz[MAXN],sum,rt,ans,k,dp[MAXN],a[MAXN],ans1;
vector<int> ed[MAXN];
void getdis(int x,int fa){
	if(!tmp[a[x]]) p[++cnt]=a[x];
	tmp[a[x]]++;
	for(int i=0;i<ed[x].size();i++){
		int y=ed[x][i];
		if(y==fa||vis[y]) continue;
		getdis(y,x);
	}
	tmp[a[x]]--;
}
bool ch[MAXN];
void dfs(int x,int fa,int r){
	if(a[x]==r) return (void)(ans++);
	for(int i=0;i<ed[x].size();i++){
		int y=ed[x][i];
		if(y==fa||vis[y]) continue;
		dfs(y,x,r);
	}
}
int num[MAXN];
int getans(int x,int y){
	int ans=0;
	cnt=0;getdis(x,0);
	for(int i=1;i<=cnt;i++) num[p[i]]++;
	for(int i=1;i<=cnt;i++){
		if(!ch[p[i]]) ans+=(num[p[i]]-1)*num[p[i]]/2;
		ch[p[i]]=1;
	}
	for(int i=1;i<=cnt;i++) ch[p[i]]=0,num[p[i]]--;
	return ans;
}
void getrt(int x,int fa){
	siz[x]=1;dp[x]=0;
	for(int i=0;i<ed[x].size();i++){
		int y=ed[x][i];
		if(y==fa||vis[y]) continue;
		getrt(y,x);
		siz[x]+=siz[y],dp[x]=max(dp[x],siz[y]);
	}
	dp[x]=max(dp[x],sum-dp[x]);
	if(dp[x]<dp[rt]) rt=x;
}
void solve(int x){
	vis[x]=1;tmp[a[x]]++;ans+=getans(x,0);tmp[a[x]]--;
	for(int i=0;i<ed[x].size();i++){
		int y=ed[x][i];
		if(vis[y]) continue;
		tmp[a[x]]++;
		ans-=getans(y,x);dfs(y,x,a[x]);
		tmp[a[x]]--;
		sum=siz[y],dp[0]=n,rt=0;
		getrt(y,x);solve(rt);
	}
}
int T;
signed main(){
	cin>>T;
	while(T--){
		ans=0;cin>>n;
		for(int i=1;i<=n;i++) ed[i].clear(),vis[i]=tmp[i]=0;
		for(int i=1;i<=n;i++) cin>>a[i];
		for(int i=1;i<n;i++){
			cin>>x>>y;
			ed[x].push_back(y);
			ed[y].push_back(x);
		}
		dp[0]=sum=n,getrt(1,0),solve(rt);
		cout<<ans<<endl;
	}
}