[gym101933]King's Colors 计数问题

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/qq_40307434/article/details/83268212

/*
    [gym101933]King's Colors
    树上染色，一共k种颜色都要用到。相邻节点颜色不同，问方案数。
*/
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
ll mod=1000000007;
const int mmax=2555;
const int nmax=2555;

int n,k;

/** 基础加边*/
int etot=0,to[mmax],nex[mmax],head[nmax];
void addedge(int u,int v){
    to[++etot]=v,nex[etot]=head[u],head[u]=etot;
}

/** vir[i]: 使用i种颜色的方案数，包括小于等于k
    rea[i]:刚刚好使用i种颜色的方案数*/
ll vir[nmax],rea[nmax];

/** dfs:当前节点rt用1种颜色方案(一共k种颜色可选)
    dfs计算vir数组。
    vir[i]=dfs(0,i)*i;
*/
ll dfs(int rt,int k){
    if(!head[rt])return 1;
    ll son_tot=1;
    for(int i=head[rt];i;i=nex[i]){
        son_tot=son_tot*(k-1)%mod*dfs(to[i],k)%mod;
        //儿子节点除了当前节点选色其他k-1种颜色都能用。
    }
    return son_tot;
}

/** 预处理计算i^(-1)%mod*/
ll ver[nmax];
ll quick_pow(ll a,ll x){
    ll ans=1;
    while(x){
        if(x&1) ans=ans*a%mod;
        a=a*a%mod;
        x>>=1;
    }return ans;
}
void init_ver(){
    for(int i=1;i<=k;i++){
        ver[i]=quick_pow(i,mod-2);
    }
}

int main(){
    int pa;
    scanf("%d%d",&n,&k);
    for(int i=1;i<n;i++){
        scanf("%d",&pa);
        addedge(pa,i);
    }
    init_ver();

    for(int i=2;i<=k;i++){
        vir[i]=i*dfs(0,i)%mod;
        rea[i]=vir[i];
        ll C=i;
        for(int j=2;j<i;j++){
            C=C*(i-j+1)%mod*ver[j]%mod;
            rea[i]=(rea[i]-C*rea[j]%mod+mod)%mod;
            //real[i]=vir[i]-sigma(2<=j<i)C(i,j)*real[j]
        }
    }
    printf("%I64d\n",rea[k]);
    return 0;
}