Description
思源湖畔有一棵树,那是独孤玉溪最喜欢的地方。
传说中,这棵不见边际的树有N个节点,每个节点都有1片叶子,每片叶子都拥有K种颜色中的一种,独孤玉溪喜欢爬到这棵树上,沿着一条路线摘叶子,并拥有所有颜色的叶子。
独孤玉溪会选择一个起点,并沿着树边走,然后最终停在一个终点上(起点和终点可能相同),当然了每一个结点只能经过一次(每一片叶子只能摘一遍)。独孤玉溪突生奇想,有多少种不同的方案能满足自己呢?(两种方案不同当且仅当起点不同或终点不同)。
Input
第一行包含两个整数N和K。
第二行包含N个整数表示col[i],为每片叶子的颜色(col[i]为1到K的一个整数)。
第三行到第N+1行,每行有两个整数x、y,表示x与y之间有一条树边。
Output
一行,表示求得的答案。
Data Constraint
20%的数据:N<=10000,K<=10
40%的数据:N<=50000,K<=2
100%的数据:N<=50000,K<=10
Solution
假如没有颜色限制,那这个n个节点的树的路径数量显然是n^2,但是有颜色限制,所以我们考虑容斥,可以先枚举那些点不被选,那么剩下的颜色显然可以选或不选,接着考虑把对应的点删去,不难发现是一个森林,暴力O(n)统计答案,最后乘上容斥系数即可
Code
#include <cstdio>
#include <cstring>
#define il inline
#define Int register int
using namespace std;
const int N = 50007;
int n, k, tot;
long long ans, sum;
int a[50007], p[11];
int first[50007], next[100007], to[100007];
bool bz[50007];
il int read()
{
int data = 0;
char ch = 0;
while (ch < '0' || ch > '9') ch = getchar();
while (ch >= '0' && ch <= '9') data = (data << 3) + (data << 1) + ch - '0', ch = getchar();
return data;
}
il void add(int x, int y)
{
next[++ tot] = first[x];
first[x] = tot;
to[tot] = y;
}
il void find(int x, int s)
{
bz[x] = true;
++ sum;
for (Int i = first[x]; i; i = next[i])
if (!bz[ to[i] ] && !(s & p[ a[ to[i] ] - 1 ])) find(to[i], s);
}
il void dfs(int x, int f, int s)
{
if (x == k)
{
memset(bz, 0, sizeof bz);
for (Int i = 1; i <= n; ++ i)
if (!bz[i] && !(s & p[a[i] - 1]))
{
sum = 0;
find(i, s);
ans += f * sum * sum;
}
return;
}
dfs(x + 1, f, s);
dfs(x + 1, -f, s + p[x]);
}
int main()
{
n = read(), k = read();
for (Int i = 1; i <= n; ++ i) a[i] = read();
for (Int i = p[0] = 1; i <= k; ++ i) p[i] = p[i - 1] << 1;
for (Int i = 1; i < n; ++ i)
{
int x = read(), y = read();
add(x, y);
add(y, x);
}
dfs(0, 1, 0);
printf("%lld", ans);
return 0;
}