Wannafly挑战赛19-C-多彩的树（状压+容斥）

链接：https://www.nowcoder.com/acm/contest/131/C
来源：牛客网
 

题目描述

有一棵树包含 N 个节点，节点编号从 1 到 N。节点总共有 K 种颜色，颜色编号从 1 到 K。第 i 个节点的颜色为 Ai。
Fi 表示恰好包含 i 种颜色的路径数量。请计算：

输入描述:

第一行输入两个正整数 N 和 K，N 表示节点个数，K 表示颜色种类数量。
第二行输入 N 个正整数，A1, A2, A3, ... ..., AN，Ai 表示第 i 个节点的颜色。

接下来 N - 1 行，第 i 行输入两个正整数 Ui 和 Vi，表示节点 Ui 和节点 Vi 之间存在一条无向边，数据保证这 N-1 条边连通了 N 个节点。

1 ≤ N ≤ 50000.
1 ≤ K ≤ 10.
1 ≤ Ai ≤ K.

输出描述:

输出一个整数表示答案。

示例1

输入

复制

5 3
1 2 1 2 3
4 2
1 3
2 1
2 5

输出

复制

4600065

#include<vector>
#include<stdio.h>
#include<string.h>
#include<algorithm>
using namespace std;
#define ll long long
#define mod 1000000007
int n,k,a[50005],flag[20],vis[50005],num[50005];
ll ans[50005],siz[50005],cnt[50005];
ll p[20];
vector<int>q[50005];
void init()
{
	p[0]=1;
	for(int i=1;i<=10;i++)
		p[i]=p[i-1]*131%mod;
	for(int i=1;i<(1<<k);i++)
		for(int j=0;j<k;j++)
			if(i&(1<<j))
				num[i]++;
}
void dfs(int x,int color)
{
	vis[x]=1;siz[x]=1;
	for(int i=0;i<q[x].size();i++)
	{
		int u=q[x][i];
		if(vis[u]==0 && (color&(1<<a[u])))
		{
			dfs(u,color);
			cnt[color]+=siz[x]*siz[u];
			siz[x]+=siz[u];
		}
	}
	cnt[color]++;
}
int main(void)
{
	int x,y;
	scanf("%d%d",&n,&k);
	init();
	for(int i=1;i<=n;i++)
		scanf("%d",&a[i]),--a[i];
	for(int i=1;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		q[x].push_back(y);
		q[y].push_back(x);
	}
	for(int i=1;i<(1<<k);i++)
	{
		memset(vis,0,sizeof(vis));
		for(int j=1;j<=n;j++)
			if(vis[j]==0 && (i&(1<<a[j])))
				dfs(j,i);
	}
	for(int i=1;i<(1<<k);i++)
	{
		for(int j=1;j<(1<<k);j++)
		{
			if(i==j) continue;
			if((i&j)==j)
				cnt[i]-=cnt[j];
		}
		ans[num[i]]+=cnt[i];
	}
	ll res=0;
	for(int i=1;i<=k;i++)
		res=(res+ans[i]*p[i]%mod)%mod;
	printf("%lld\n",res);
	return 0;
}