【GDOI2017模拟一试4.11】颜色树

明明这么简单的题目调的我死去活来的。。
题意：给你一棵树，树上的点都有一个颜色的属性，然后，问你找一条路径，能经过所有颜色，有多少种方案。

分析：xjb容斥一波，暴力枚举那些点不被选择，然后dfs一遍，求出当前图的路径数量，乘以容斥系数累加起来就好了。

#include<cstdio> 
#include<cstring>
#include<algorithm>
#include<iostream>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
using namespace std;
int n,m,k;
typedef long long ll;
const int N=3e5+5;
int head[N],next[N],go[N],tot;
int bz[15];
int col[N],vis[50005];
ll ans=0,cnt;
inline int dfs1(int x)
{
    vis[x]=1;
    ++cnt;
    for(int i=head[x];i;i=next[i])
    {
        int v=go[i];
        if (!vis[v]&&!bz[col[v]])
        {
            dfs1(v);
        }
    }
}
inline void dfs(int x,int bo,int num)
{
    if (x==m+1)return;
    dfs(x+1,0,num);
    bz[x+1]=1;
    dfs(x+1,1,num+1);
    bz[x+1]=0;
    if (x==m)
    {
        memset(vis,0,sizeof(vis));
        ll sum=0;
        fo(i,1,n)if (!vis[i]&&!bz[col[i]])
        {
            cnt=0;
            dfs1(i);
            sum+=(ll)cnt*cnt;
        }
        if (num%2==0)
        ans+=sum;
        else ans-=sum;
    }
}
inline void add(int x,int y)
{
    go[++tot]=y;
    next[tot]=head[x];
    head[x]=tot;
}
int main()
{
    freopen("colortree.in","r",stdin);
    freopen("colortree.out","w",stdout);
    scanf("%d%d",&n,&m);
    fo(i,1,n)scanf("%d",&col[i]);
    fo(i,1,n-1)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x);
    }
    k=1;
    ans=0;
    dfs(1,0,0);
    bz[1]=1;
    dfs(1,1,1);
    //dfs(1,1,1); 
    printf("%lld\n",ans);
}