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 21 2 21 21 3

Sample Output

6

Data Constraint

20%的数据：N<=10000,K<=1040%的数据：N<=50000,K<=2100%的数据：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 51010int 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;}