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;elsedp2[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;
}