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

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

题目描述

有一棵树包含 N 个节点，节点编号从 1 到 N。节点总共有 K 种颜色，颜色编号从 1 到 K。第 i 个节点的颜色为 A _i。
F _i 表示恰好包含 i 种颜色的路径数量。请计算：
$\left( \sum _{ i=1 }^{ K }{ ({ F }_{ i } \times { 131 }^{ i }) } \right) { mod }({ 10 }^{ 9 }+7)$

输入描述:

第一行输入两个正整数 N 和 K，N 表示节点个数，K 表示颜色种类数量。
第二行输入 N 个正整数，A₁, A₂, A₃, ... ..., A_N，A_i 表示第 i 个节点的颜色。

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

1 ≤ N ≤ 50000. 1 ≤ K ≤ 10. 1 ≤ A _i ≤ K.

输出描述:

输出一个整数表示答案。

因为k只有10，所以可以暴力所有2^k种情况，对于每一种情况计算有多少种不同的路径

设dp[x]表示x状态下路径的条数，例如假设k=5，x=10010(二进制)，那么dp[x]就相当在保留所有颜色为2,5的节点，删除所有其它颜色节点的情况下路径的条数

不过因为dp[x]并不是经过当前所有保留颜色的路径条数，所以还要容斥一下，直接2^20暴力即可，奇加偶减，这样就可以求出总共经过i种不同颜色的路径条数了

#include<stdio.h>
#include<string.h>
#include<vector>
using namespace std;
#define LL long long
#define mod 1000000007
int col[50005], p[55], vis[50005];
LL sum, ans[15], cds[15] = {1}, dp[1111], num[1111], dp2[1111];
vector<int> G[50005];
void Sech(int x)
{
	int i, v;
	vis[x] = 1, sum++;
	for(i=0;i<G[x].size();i++)
	{
		v = G[x][i];
		if(vis[v] || p[col[v]]==0)
			continue;
		Sech(v);
	}
}
int main(void)
{
	LL cot;
	int n, k, i, j, x, y;
	scanf("%d%d", &n, &k);
	for(i=1;i<=k;i++)
		cds[i] = cds[i-1]*131%mod;
	for(i=1;i<=n;i++)
		scanf("%d", &col[i]);
	for(i=1;i<=n-1;i++)
	{
		scanf("%d%d", &x, &y);
		G[x].push_back(y);
		G[y].push_back(x);
	}
	for(i=1;i<(1<<k);i++)
	{
		x = 0;
		for(j=0;j<=k-1;j++)
		{
			p[j+1] = 0;
			if(i&(1<<j))
				p[j+1] = 1, x++;
		}
		for(j=1;j<=n;j++)
			vis[j] = 0;
		for(j=1;j<=n;j++)
		{
			if(vis[j]==0 && p[col[j]])
			{
				sum = 0;
				Sech(j);
				dp2[i] = dp[i] = (dp[i]+sum*(sum-1)/2)%mod;
				num[i] = x;
			}
		}
	}
	for(i=1;i<(1<<k);i++)
	{
		for(j=1;j<(1<<k);j++)
		{
			if((i|j)==i && i!=j)
			{
				if((num[i]-num[j])%2)
					dp2[i] = (dp2[i]-dp[j]+mod)%mod;
				else
					dp2[i] = (dp2[i]+dp[j])%mod;
			}
		}
		ans[num[i]] = (ans[num[i]]+dp2[i])%mod;
	}
	cot = 0;
	for(i=1;i<=k;i++)
		cot = (cot+ans[i]*cds[i])%mod;
	cot = (cot+n*131)%mod;
	printf("%lld\n", cot);
	return 0;
}