D. Petya and Coloring 题解(组合数学+dp+容斥)

题目链接

构造方法比较容易想到

第1列和第m列的颜色种数要相等，中间的列颜色来源于第1列和第m列中的共同颜色

主要的问题是如何解决\(n\)个元素里面存在\(i\)个元素，且每个元素至少存在一次的方案数

这是一个经典问题，可以利用dp+容斥去解决

代码

#include<bits/stdc++.h>
#define fi first
#define se second
#define debug cout<<"I AM HERE"<<endl;
using namespace std;
typedef long long ll;
const int maxn=1e6+5,inf=0x3f3f3f3f,mod=1e9+7;
const double eps=1e-6;
int n,m,k;
ll fac[maxn],finv[maxn];
ll dp[1000+5];
ll qpow(ll a,ll b){
    ll ans=1,base=a;
    while(b){
        if(b&1) ans=ans*base%mod;
        base=base*base%mod;
        b=b>>1;
    }
    return ans;
}
void init(){
    fac[0]=finv[0]=1;
    for(int i=1;i<=1e6;i++){
        fac[i]=fac[i-1]*i%mod;
    }
    finv[1000000]=qpow(fac[1000000],mod-2);
    for(int i=1e6-1;i>=1;i--){
        finv[i]=finv[i+1]*(i+1)%mod;
    }
}
ll c(ll a,ll b){
    if(a<b) return 0;
    ll ans=fac[a]*finv[b]%mod*finv[a-b]%mod;
    return ans;
}
signed main(){
    init();
    cin>>n>>m>>k;
    if(m==1){
        printf("%lld\n",qpow(k,n));
        return 0;
    }
    for(int i=1;i<=n;i++){
        dp[i]=qpow(i,n)%mod;
        for(int j=1;j<i;j++){
            dp[i]=((dp[i]-dp[j]*c(i,j))%mod+mod)%mod;
        }
    }
    ll pr=0;
    int beg=0;
    for(int i=beg;i<=min(n,k);i++){
        for(int j=0;j<=min(n,k);j++){
            if(i+j>min(k,n)) continue;
            ll ans=dp[i+j]*dp[i+j]%mod;
            ans=ans*c(k,i)%mod*c(k-i,j)%mod*c(k-i-j,j)%mod*qpow(i,(m-2)*n)%mod;
            pr=(pr+ans)%mod;
        }
    }
    printf("%lld\n",pr);
    return 0;
}