链接:https://www.nowcoder.com/acm/contest/131/C
来源:牛客网
题目描述
有一棵树包含 N 个节点,节点编号从 1 到 N。节点总共有 K 种颜色,颜色编号从 1 到 K。第 i 个节点的颜色为 A
i。
F i 表示恰好包含 i 种颜色的路径数量。请计算:
F i 表示恰好包含 i 种颜色的路径数量。请计算:
输入描述:
第一行输入两个正整数 N 和 K,N 表示节点个数,K 表示颜色种类数量。 第二行输入 N 个正整数,A1, A2, A3, ... ..., AN,Ai 表示第 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;
}