jzoj 5050 颜色树

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

Time Limits: 2000 ms

---

解题思路

显然，枚举每条路径是否合法，会超时。
发现k只有10，可以使用容斥原理
如果去掉颜色的限制，那么对于一棵 n 个节点的树，它的路径个数就是 n^2 了。
整个图就k种颜色，所以ans=n^2-颜色1不选-颜色2不选+颜色1,2选……，以此类推，于是ans=n^2-任意1种颜色不选+任意2种颜色不选……
枚举哪些颜色不被选，剩下其它颜色可以选或不选。然后把这些颜色对应的点全部删去，得到一个森林，然后计算出这种情况路径种数，最后乘容斥系数就可以了。

---

#include<iostream>
#include<cstring>
#include<cstdio>
#include<algorithm>
#include<cmath>
using namespace std;

#define K 2500
#define N 51010

int s[K],o[K];

struct note
{
    int x,y,ne;
};

note side[N*2];
int co[N],last[N],p[N];
int n,kk,l,t,ans;

void add(int x,int y)
{
    ++l;
    side[l].x=x;
    side[l].y=y;
    side[l].ne=last[x];
    last[x]=l;
}

void addside(int x,int y)
{
    add(x,y);
    add(y,x);
}

void init()
{
    scanf("%d%d",&n,&kk);
    for (int i=1;i<=n;++i)
      scanf("%d",&co[i]);
    memset(last,0,sizeof(last));
    for (int i=1;i<=n-1;++i)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        addside(x,y);
    }
}

void dfs(int x,int k)
{
    t++;
    p[x]=1;
    //printf("%d",x);
    for (int i=last[x];i!=0;i=side[i].ne)
    {
        int j=side[i].y;
        if (((1<<(co[j]-1))&k)!=0) continue;
        if (p[j]==1) continue;
        dfs(j,k);
    }
}

void gets()
{
    for (int k=0;k<=(1<<kk)-1;++k)
    {
        s[k]=0;
        memset(p,0,sizeof(p));
        for (int i=1;i<=n;++i)
          if (p[i]==0 && (((1<<(co[i]-1))&k)==0))
            { 
                t=0;
            dfs(i,k);
            s[k]+=t*t;
          }
    }
    for (int k=0;k<=(1<<kk)-1;++k)
    {
        int i=k;
        o[k]=0;
        while (i>0)
        {
            o[k]+=i%2;
            i/=2;
        }
    }
}

int main()
{
    freopen("colortree.in","r",stdin);
    freopen("colortree.out","w",stdout);
    init();
    gets();
    ans=0;
    for (int k=0;k<=(1<<kk)-1;++k)
    {
      if (o[k]%2==1) ans-=s[k];
        else ans+=s[k];
    }
    printf("%d\n",ans);
    return 0;
}