CF1923E

#include<iostream>
#include<algorithm>
#include<map>
#define ll int
using namespace std;
const ll maxn=2e5+50;
ll a[maxn];
ll hed[maxn];
ll tot=0,ans;
map<ll,ll> num[maxn],h[maxn];
struct node{
	ll to,nxt;
}edge[maxn];
void add(ll u,ll v){
	edge[++tot]=(node){v,hed[u]};
	hed[u]=tot;
}
void dfs(ll x,ll fa){
	num[x][a[x]]=1;
	for(ll i=hed[x];i;i=edge[i].nxt){
		ll v=edge[i].to;
		if(v==fa) continue;
		dfs(v,x);
		ans+=num[v][a[x]];
		for(auto [col,nm]:num[v]){
			if(col==a[x]) continue;
			ans+=num[x][col]*num[v][col];
			num[x][col]+=nm;
		} 
	}
	num[x][a[x]]=1;
}
void solve(){
	ll n;
	cin>>n;
	tot=0,ans=0;
	for(ll i=1;i<=n;i++) cin>>a[i],hed[i]=0,num[i].clear();
	for(int i=1;i<n;i++){
		ll u,v;
		cin>>u>>v;
		add(u,v),add(v,u);
	}
	dfs(1,0);
	cout<<ans<<endl;
}
int main(){
	ll t=1;
	cin>>t;
	while(t--){
		solve();
	}
	return 0;
}

求助大佬，这个题为什么tle