NC17061 多彩的树

状态表示不难想出，可以通过dfs搜索连通块，枚举每一种情况对应的路径数。由于在计数过程中有重复，需要对每个dp[i]减去 情况i包含所有子情况j对应的dp值。 几个重点：

　　1、枚举子情况的循环，需要牢记。

　　2、在dp[i]-=dp[j]后，有可能出现dp[i]<0，进而导致ans<0，故类似问题尽量写成ans=(ans+mod)%mod。

#include<bits/stdc++.h>
using namespace std;
#define rep(i,l,r) for(int (i)=(l);(i)<=(r);++(i))
#define drep(i,r,l) for(int (i)=(r);(i)>=(l);--(i))
#define ll long long
const double eps=1e-8;
inline char nc(){
    static char buf[100000],*p1=buf,*p2=buf;
    return p1==p2&&(p2=(p1=buf)+fread(buf,1,100000,stdin),p1==p2)?EOF:*p1++;
}
#ifndef ONLINE_JUDGE
#define nc() getchar()
#endif
#define ll long long
template<typename T>
inline T read(){
    T x=0,y=1;
    char ch=nc();
    while(ch<'0'||ch>'9'){
        if(ch=='-')y=-y;
        ch=nc();
    }
    while(ch>='0'&&ch<='9'){
        x=10*x+(ch^48);
        ch=nc();
    }
    return x*y;
}
#define readl() read<ll>()
#define readi() read<int>()
#define maxn 5000200
#define ll long long
const ll mod=1e9+7;
int n,k;
ll a[maxn];
struct node{
    int to,next;
}e[maxn];
int head[maxn],tot;
bool used[maxn];
ll dp[maxn],cnt;
ll mul[maxn];
inline void add(int x, int y){
    e[++tot].to=y;
    e[tot].next=head[x];
    head[x]=tot;
}
void dfs(int u,int cur,int from){//control the current sequence to a stable value. (this can be inclusive)
    if((cur|a[u])!=cur||used[u])return;
    used[u]=1;
    ++cnt;
    for(int i=head[u];i;i=e[i].next){
        if(e[i].to!=from)dfs(e[i].to,cur,u);
    }
}
signed main(){
    n=readi(),k=readi();
    mul[0]=1;
    rep(i,1,n){
        ll x=readl();
        a[i]=1<<(x-1);
        mul[i]=(mul[i-1]*131)%mod;
    }
    rep(i,1,n-1){
        int x=readi(),y=readi();
        add(x,y);add(y,x);
    }
    rep(i,1,(1<<k)-1){
        memset(used,0,sizeof(used));
        rep(j,1,n){
            if(!used[j]){
                cnt=0;
                dfs(j,i,0);
                dp[i]=(dp[i]+cnt+cnt*(cnt-1)/2)%mod;
            }
        }
    }
    ll ans=0;
    rep(i,1,(1<<k)-1){
        for(int j=(i-1)&i;j;j=(j-1)&i){
            dp[i]-=dp[j];
        }
        ans=(ans%mod+1ll*dp[i]*mul[__builtin_popcount(i)]%mod)%mod;
    }
    printf("%lld\n",(ans+mod)%mod);
    return 0;
}

View Code