JZOJ5050 [GDOI2017模拟一试4.11] 颜色树

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