Wannafly挑战赛19-C-多彩的树（状压+容斥）

链接：https://www.nowcoder.com/acm/contest/131/C
来源：牛客网

多彩的树

时间限制：C/C++ 5秒，其他语言10秒
空间限制：C/C++ 262144K，其他语言524288K
64bit IO Format: %lld

题目描述

有一棵树包含 N 个节点，节点编号从 1 到 N。节点总共有 K 种颜色，颜色编号从 1 到 K。第 i 个节点的颜色为 Ai。
Fi 表示恰好包含 i 种颜色的路径数量。请计算：

输入描述:

第一行输入两个正整数 N 和 K，N 表示节点个数，K 表示颜色种类数量。
第二行输入 N 个正整数，A1, A2, A3, ... ..., AN，Ai 表示第 i 个节点的颜色。

接下来 N - 1 行，第 i 行输入两个正整数 Ui 和 Vi，表示节点 Ui 和节点 Vi 之间存在一条无向边，数据保证这 N-1 条边连通了 N 个节点。

1 ≤ N ≤ 50000.
1 ≤ K ≤ 10.
1 ≤ Ai ≤ K.

输出描述:

输出一个整数表示答案。

示例1

输入

5 3
1 2 1 2 3
4 2
1 3
2 1
2 5

输出

4600065

思路：对于一个大小为x的树，它的路径条数为x+x*(x-1)。

由于k<=10，考虑直接枚举当前所要计算的颜色的状态zt，然后直接遍历所有点，如果点u的颜色包含在这个状态里，就记jd[u]=1并且累加其所有儿子节点v的jd[v]，否则记jd[u]=0。

这样，我们就可以依次查找每个点，如果u被包含在这个状态里，但是它的父节点f[u]不包含在这个状态里，那么以u为根的子树就形成了一个独立的全部包含于zt的连通块。这时候就可以累加它的路径条数。

然鹅这时候我们求的是 包含颜色 <= zt中1的个数 的路径条数。因此求完以后，我们再容斥一下，减去 包含颜色 <= zt中1的个数-1 的路径条数即可。由于不影响最终答案，因此含有相同1个数的zt我们可以分开计算，累加求和。

代码：

#include<bits/stdc++.h>
#define ll long long
#define inf 0x3f3f3f3f
#define rep(i,a,b) for(register int i=(a);i<=(b);i++)
#define dep(i,a,b) for(register int i=(a);i>=(b);i--)
using namespace std;
const int maxn=1e5+5;
//const double pi=acos(-1.0);
//const double eps=1e-9;
const ll mo=1e9+7;
int n,m,k;
int a[maxn],b[maxn],c[maxn];
int flag,cnt,num,cm;
int ok[maxn],co[maxn],jd[maxn];
int he[maxn],fa[maxn];
ll f[maxn],pw[maxn];
struct node{
    int v,nxt;
}e[maxn<<1];
template <typename T>
inline void read(T &X)
{
    X=0;int w=0; char ch=0;
    while(!isdigit(ch)) {w|=ch=='-';ch=getchar();}
    while(isdigit(ch)) X=(X<<3)+(X<<1)+(ch^48),ch=getchar();
    if(w) X=-X;
}
void add(int u,int v){
    e[cnt].v=v;
    e[cnt].nxt=he[u];
    he[u]=cnt++;
}
int count(int x){
    int sum=0;
    while(x){
        x&=(x-1);
        sum++;
    }
    return sum;
}
bool cmp(int x,int y){return c[x]<c[y];}
void dfs(int u,int ff,int zt){
    fa[u]=ff;
    if(zt&(1<<(a[u]-1))) jd[u]=1;
    else jd[u]=0;
    for(int i=he[u];i;i=e[i].nxt){
        int v=e[i].v;
        if(v==ff) continue;
        dfs(v,u,zt);
        if(jd[u]) jd[u]+=jd[v];
    }
}
int main()
{
    pw[0]=1;
    rep(i,1,10) pw[i]=(pw[i-1]*131)%mo;
    int T,cas=1;
    //read(T);
    //while(T--)
    {
        read(n);read(k);
        cnt=1;
        rep(i,1,n) read(a[i]);
        rep(i,1,n-1) {
            int u,v;
            read(u);read(v);
            add(u,v);
            add(v,u);
        }
        int m=(1<<k)-1;
        rep(i,1,m) {c[i]=count(i);}
        ll ans=0;
        rep(s,1,m){
            f[s]=0;
            dfs(1,0,s);
            rep(i,1,n) if(jd[i]&&!jd[fa[i]]){
                f[s]+=jd[i]+(ll)jd[i]*(jd[i]-1)/2;
            }
        }
        rep(s,1,m){
            for(int sub=s&(s-1);sub;sub=(sub-1)&s){
                f[s]-=f[sub];
            }
            ans=(ans+f[s]*pw[c[s]]%mo)%mo;
        }
        printf("%lld\n",ans);
    }
    return 0;
}