JZOJ 5050. 【GDOI2017模拟一试4.11】颜色树

Description

思源湖畔有一棵树，那是独孤玉溪最喜欢的地方。
传说中，这棵不见边际的树有N个节点，每个节点都有1片叶子，每片叶子都拥有K种颜色中的一种，独孤玉溪喜欢爬到这棵树上，沿着一条路线摘叶子，并拥有所有颜色的叶子。
独孤玉溪会选择一个起点，并沿着树边走，然后最终停在一个终点上(起点和终点可能相同)，当然了每一个结点只能经过一次(每一片叶子只能摘一遍)。独孤玉溪突生奇想，有多少种不同的方案能满足自己呢？(两种方案不同当且仅当起点不同或终点不同)。

Input

第一行包含两个整数N和K。
第二行包含N个整数表示col[i]，为每片叶子的颜色(col[i]为1到K的一个整数)。
第三行到第N+1行，每行有两个整数x、y，表示x与y之间有一条树边。

Output

一行，表示求得的答案。

Sample Input

3 2
1 2 2
1 2
1 3

Sample Output

6

Data Constraint

20%的数据：N<=10000,K<=10
40%的数据：N<=50000,K<=2
100%的数据：N<=50000，K<=10

分析

正解点分治根本调不出。。调到生无可恋
如果去掉颜色的限制，那么对于一棵n个节点的树，它的路径个数就是n^2了。
现在有了颜色的限制，直接算不好算，但是k不大，可以考虑容斥。
枚举哪些颜色不被选，剩下其它颜色可以选或不选。然后把这些颜色对应的点全部删去，
得到一个森林，接下来O(n)统计答案，最后乘上容斥系数即可。只要会容斥原理，这
种方法是没什么理解难度的。
时间复杂度O(2^k*n)

代码

#include <bits/stdc++.h>

#define N 50005
#define ll long long

struct NOTE
{
    int to,next;
}e[N * 2];

int cnt;
int next[N];

int n,m;

int color[N];
bool vis[N];

int size;
ll ans;

void add(int x,int y)
{
    e[++cnt].to = y;
    e[cnt].next = next[x];
    next[x] = cnt;
}

void getAns(int fa,int s)
{
    vis[fa] = 1;
    size++;
    for (int i = next[fa]; i; i = e[i].next)
    {
        if (!vis[e[i].to] && (s & color[e[i].to]) == 0)
            getAns(e[i].to,s);
    }
}

void dfs(int fa,int f,int s)
{
    if (fa > m)
    {
        memset(vis,0,sizeof(vis));
        for (int i = 1; i <= n; i++)
            if (!vis[i] && (color[i] & s) == 0)
            {
                size = 0;
                getAns(i,s);
                ans += (ll) size * size * f;
            }
        return;
    }
    dfs(fa + 1,f,s);
    dfs(fa + 1, -f, s + (1 << (fa - 1)));
}

int main()
{
    freopen("colortree.in","r",stdin);
    freopen("colortree.out","w",stdout);
    scanf("%d%d",&n,&m);
    for (int i = 1; i <= n; i++)
    {
        int x;
        scanf("%d",&x);
        color[i] = 1 << (x - 1);
    }
    for (int i = 1; i < n; i++)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x);
    }

    dfs(1,1,0);

    printf("%lld\n",ans);
}