HDU - 5977

题目链接：HDU - 5977

---

分治重心，然后对根经过的不同颜色之和状态压缩。

之后容斥统计答案。

---

AC代码：

#pragma GCC optimize("-Ofast","-funroll-all-loops")
#include
#define int long long
using namespace std;
const int N=5e4+10;
int n,rt,col[N],k,sz[N],mx,sum,vis[N],res,st[1<<10],q[N*5],d[N],tots;
vector<int> g[N];
inline void add(int a,int b){g[a].push_back(b),g[b].push_back(a);}
void get(int x,int fa){
	sz[x]=1;	int t=0;
	for(auto to:g[x]){
		if(vis[to]||to==fa)	continue;
		get(to,x);	sz[x]+=sz[to];	t=max(t,sz[to]);
	} 
	t=max(t,sum-sz[x]); if(t<mx) mx=t,rt=x;
}
void getdis(int x,int fa,int s){
	q[++q[0]]=s;	st[s]++;
	for(auto to:g[x]){
		if(to==fa||vis[to])	continue;
		getdis(to,x,s|(1<<col[to]));
	}
}
inline int calc(int x,int s){
	memset(st,0,sizeof st);
	q[0]=0;	int res=0;	getdis(x,0,s|(1<<col[x]));
	for(int i=0;i<k;i++){
		for(int s=tots;s>=0;s--) if(!(s>>i&1))	st[s]+=st[s|(1<<i)];
	}
	for(int i=1;i<=q[0];i++)	res+=st[tots^q[i]];
	return res;
}
void divide(int x){
	vis[x]=1;	res+=calc(x,0);
	for(auto to:g[x]){
		if(vis[to])	continue;	res-=calc(to,1<<col[x]);
		mx=sum=sz[to]; get(to,x); get(rt,0);
		divide(rt);
	}
}
inline void solve(){
	for(int i=1;i<=n;i++)	scanf("%lld",&col[i]),col[i]--,g[i].clear(),vis[i]=0;
	tots=(1<<k)-1;	res=0;
	for(int i=1,a,b;i<n;i++)	scanf("%lld %lld",&a,&b),add(a,b);
	if(k==1)	return printf("%lld\n",n*n),void();
	mx=sum=n;	get(1,0);	get(rt,0);
	divide(rt);
	printf("%lld\n",res);
}
signed main(){
	while(cin>>n>>k)	solve();
	return 0;
}