【HAOI2018】染色

纪念独立想出的第一个生成函数题！题目链接

题意:
有 \(n\) 个位置和 \(m\) 种颜色，给出一个数 \(s\)，若恰好有 \(k\) 种颜色出现 \(s\) 次，则该染色方案的权值为 \(W_k\)。求所有染色方案的权值和。

我们不妨对每个 \(k\) 求方案数 \(C_k\)，最终答案为 \(\sum C_kW_k\)。

这种至多/至少转恰好的形式，大概是二项式反演。略加思考，发现至少比较好求。设钦定 \(k\) 个颜色出现 \(s\) 次的方案数为 \(g_k\) （也就是至少），恰好有 \(k\) 个的方案数为 \(f_k\)。

\(g_k\) 比较好求，有

\[g_k=\binom{m}{k}\binom{n}{ks}\binom{ks}{s,\cdots,s}(m-k)^{n-ks}=\binom{m}{k}\frac{n!}{(s!)^k(n-ks)!}(m-k)^{n-ks} \]

考虑到每一个 \(g_k\) 钦定 \(k\) 种颜色时，\(f_j(j\geq k)\) 被重复计算了 \(\binom{j}{k}\) 次。（因为钦定 \(k\) 个颜色时，不同的钦定方案可能会计算出相同的结果）

于是

\[g_k=\sum_{i=k}^m \binom{i}{k} f_i \]

由二项式反演得

\[\begin{aligned} f_k&=\sum_{i=k}^m (-1)^{i-k}\binom{i}{k}g_i \\ &=\sum_{i=k}^m (-1)^{i-k}\frac{i!}{(i-k)!k!}g_i \\ \end{aligned} \]

设 \(F_i=\frac{(-1)^i}{i!}\)， \(G_i=i!g_i\) 。
可得

\[f_k=\frac{1}{k!}\sum_{i=k}^m F_{i-k} G_{i} \]

不妨翻转一下 \(G\) 使得 \(G_i=G_{m-i}\)

\[\begin{aligned} f_k&=\frac{1}{k!}\sum_{i=k}^m F_{i-k} G_{m-i}\\ &=\frac{1}{k!}\sum_{i=0}^{m-k} F_{i} G_{m-k-i} \end{aligned} \]

然后NTT卷积一下就可以了QwQ

#include <bits/stdc++.h>
#define ll long long
using namespace std;

const int maxn=5e5+5;
const int maxnum=1e7+5;
const ll mod=1004535809;

ll qpow(ll a,ll b){
	if(b<0) return 0;
    ll ret=1;
    while(b){
        if(b&1) ret=ret*a%mod;a=a*a%mod;b>>=1;
    }
    return ret;
}

namespace NTT{
const ll g=3;
const ll gi=qpow(g,mod-2);
int t[maxn];
void init(int n){
    for(int i=1;i<n;i++) t[i]=(t[i>>1]>>1)|(i&1?n>>1:0);
}
void NTT(ll *y,int n,int inv){
    for(int i=0;i<n;i++) if(i<t[i]) std::swap(y[t[i]],y[i]);
    for(int len=2;len<=n;len<<=1){
        ll omg=qpow(inv==1?g:gi,(mod-1)/len);
        for(int i=0;i<n;i+=len){
            ll cur=1;
            for(int j=i;j<i+len/2;j++){
                ll e=y[j];ll o=cur*y[j+len/2]%mod;
                y[j]=(e+o)%mod;y[j+len/2]=(e-o+mod)%mod;
                cur=cur*omg%mod;
            }
        }
    }
    if(inv==-1){
        ll div=qpow(n,mod-2);
        for(int i=0;i<n;i++) y[i]=y[i]*div%mod;
    }
}
}

int get(int x){
    if(x&(-x)==x) return x;
    int k=x&(-x);while(k<x) k<<=1;
    return k;
}

int n,m,s;
ll fac[maxnum],w[maxn],ans;
ll g[maxn];
ll G[maxn],F[maxn];

ll C(ll n,ll m){
	return fac[n]*qpow(fac[m],mod-2)%mod*qpow(fac[n-m],mod-2)%mod;
}

int main(){
	scanf("%d%d%d",&n,&m,&s);
	fac[0]=1;
	for(int i=0;i<=m;i++) scanf("%lld",&w[i]);
	for(int i=1;i<maxnum;i++) fac[i]=fac[i-1]*i%mod;
	for(int i=0;i<=m;i++){
		if(n-i*s<0) break;
		g[i]=C(m,i)*fac[n]%mod*qpow(m-i,n-i*s)%mod*qpow(qpow(fac[s],i),mod-2)%mod*qpow(fac[n-i*s],mod-2)%mod;
	}
	int cur=1;
	for(int i=0;i<=m;i++){
		G[m-i]=fac[i]*g[i]%mod;
		if(cur>0){
			F[i]=qpow(fac[i],mod-2);
		}else{
			F[i]=mod-qpow(fac[i],mod-2);
		}
		cur=-cur;
	}
	int len=get(2*m+1);
	NTT::init(len);
	NTT::NTT(F,len,1),NTT::NTT(G,len,1);
	for(int i=0;i<len;i++) F[i]=F[i]*G[i]%mod;
	NTT::NTT(F,len,-1);
	for(int i=0;i<=m;i++){
		ans+=w[i]*qpow(fac[i],mod-2)%mod*F[m-i]%mod;
		if(ans>mod) ans-=mod;
	}
	printf("%lld",ans);
}