LOJ2527「HAOI2018」染色

Description

题目链接

给定长为 \(n\) 的序列，你需要用 \(m\) 种颜色为其染色。若一种染色方案中恰好出现 \(S\) 次的颜色有 \(K\) 种，那么它的代价为 \(W_K\)。求所有可能的染色方案的代价总和对 \(1004535809\) 取模的结果

\(n\le 10^7,m\le 10^5,S\le 150,0\le W_i<1004535809\)

Solution

考虑到如果一种染色方案中出现 \(S\) 次的颜色种数定了，它的代价也就定了，所以我们可以先算出恰好出现 \(S\) 次的颜色有 \(K\) 种的方案数，再对应乘上 \(W_K\) 就是我们要的结果了

先写出暴力的 \(\text{DP}\) 方程，设 \(dp_{i,j,k}\) 表示已经用前 \(i\) 种颜色给序列上 \(j\) 个位置染过色了，其中恰好染了 \(S\) 次的颜色有 \(k\) 种的方案数，那么有

\[dp_{i,j,k}=\dbinom{n-j+S}{S}dp_{i-1,j-S,k-1}+\sum\limits_{0\le c\le j且c\not=S}\dbinom{n-j+c}{c}dp_{i-1,j-c,k} \]

答案就是 \(\sum\limits_{i=0}^{m}dp_{m,n,i}W_i\)，复杂度 \(O(n^2m^2)\)

考虑生成函数？状态数太多

发现没法优化，仅仅是状态数就已经是 \(O(nm^2)\) 了

回到题目，我们要求的是出现 \(S\) 次的颜色恰好有 \(K\) 种的方案数，这样才能方便我们算出代价

恰好？考虑广义容斥原理

其实和一般容斥原理也没什么区别，就在于系数的问题

组合数形式的容斥原理中，一个具有 \(K\) 个性质的方案会在我们限制至少需要满足 \(i\) 种性质时被计算到 \(\binom{K}{i}\) 次，而我们想要求得一组 \(\{f_n\}\)，使得

\[[x==K]=\sum\limits_{i=0}^{x}\dbinom{x}{i}f_i \]

二项式反演得到

\[f_n=\sum\limits_{i=0}^{n}(-1)^{n-i}\dbinom{n}{i}[i==K] \]

即

\[f_n=(-1)^{n-K}\dbinom{n}{K} \]

回到题目，现在我们已经得到了容斥系数，考虑计算限制出现 \(S\) 次的颜色至少有 \(i\) 种的方案数，设这一步的方案数为 \(g_i\)，那么有

\[g_i=\dbinom{m}{i}\frac{n!}{(S!)^i(n-iS)!}(m-i)^{n-iS} \]

现在求出现 \(S\) 次的颜色恰好有 \(K\) 种的方案数就不成问题了

设 \(f_{i,j}=(-1)^{i-j}\dbinom{i}{j}\)，\(c_i\) 表示出现 \(S\) 次的颜色恰好有 \(i\) 种的方案数，那么有

\[c_K=\sum\limits_{i=0}^{m}f_{i,K}g_i \]

直接代入 \(f\) 与 \(g\) 的值，得

\[c_K=\sum\limits_{i=0}^{m}(-1)^{i-K}\dbinom{i}{K}\dbinom{m}{i}\frac{n!}{(S!)^i(n-iS)!}(m-i)^{n-iS} \]

先简单地换个元

\[c_i=\sum\limits_{j=0}^{m}(-1)^{j-i}\dbinom{j}{i}\dbinom{m}{j}\frac{n!}{(S!)^j(n-jS)!}(m-j)^{n-jS} \]

把组合数暴力展开，得

\[c_i=\frac{m!n!}{i!}\sum\limits_{j=0}^{m}(-1)^{j-i}\frac{1}{(j-i)!}\frac{1}{(m-j)!}\frac{1}{(S!)^j(n-jS)!}(m-j)^{n-jS} \]

重定义 \(f_i=(-1)^i\frac{1}{i!}\) ，\(g_j=\frac{1}{(m-j)!}\frac{1}{(S!)^j(n-jS)!}(m-j)^{n-jS}\)

那么

\[c_i=\frac{m!n!}{i!}\sum\limits_{j=i}^{m}f_{j-i}g_j \]

将 \(f\) 翻转，得 \(f'\) ，那么

\[c_i=\frac{m!n!}{i!}\sum\limits_{j=i}^{m}f'_{m+i-j}g_j \]

直接可以写成

\[c_i=\frac{m!n!}{i!}\sum\limits_{j=0}^{m+i}f'_{m+i-j}g_j \]

\(\text{NTT}\) 即可，复杂度 \(O(n+m\log m)\)

代码如下：

#include<cstdio>
#include<iostream>
#include<algorithm>
using namespace std;
const int N=1e5+10;
const int M=1e7+10;
const int mod=1004535809;
const int G=3;
const int invG=334845270;
int n,m,s,w[N],fac[M],inv[M],A,d[N],f[N<<2],g[N<<2],k,INV,c[N],ans;
inline void Add(int &x,int y){x+=y;x-=x>=mod? mod:0;}
inline int MOD(int x){x-=x>=mod? mod:0;return x;}
inline int Minus(int x){x+=x<0? mod:0;return x;}
inline int fas(int x,int p){int res=1;while(p){if(p&1)res=1ll*res*x%mod;p>>=1;x=1ll*x*x%mod;}return res;}
inline void Preprocess(){
	int t=max(n,max(m,s));
	fac[0]=1;for(register int i=1;i<=t;i++)fac[i]=1ll*fac[i-1]*i%mod;
	inv[t]=fas(fac[t],mod-2);inv[0]=inv[1]=1;
	for(register int i=t-1;i>=2;i--)inv[i]=1ll*inv[i+1]*(i+1)%mod;
}
inline void NTT(int *a,int f){
	for(register int i=0,j=0;i<k;i++){
		if(i>j)swap(a[i],a[j]);
		for(register int l=k>>1;(j^=l)<l;l>>=1);}
	for(register int i=1;i<k;i<<=1){
		int w=fas(~f? G:invG,(mod-1)/(i<<1));
		for(register int j=0;j<k;j+=(i<<1)){
			int e=1;
			for(register int p=0;p<i;p++,e=1ll*e*w%mod){
				int x=a[j+p],y=1ll*a[j+p+i]*e%mod;
				a[j+p]=MOD(x+y);a[j+p+i]=MOD(x-y+mod);
			}
		}
	}
}
int main(){
	scanf("%d%d%d",&n,&m,&s);Preprocess();
	for(register int i=0;i<=m;i++)scanf("%d",&w[i]);
	A=1ll*fac[n]*fac[m]%mod;
	for(register int i=0;i<=m;i++)d[i]=1ll*A*inv[i]%mod;
	for(register int i=0;i<=m;i++)
		f[i]=1ll*((i&1)? mod-1:1)*inv[i]%mod;
	for(register int i=0;i<=m;i++)
		if(n-i*s>=0)g[i]=1ll*inv[m-i]*fas(inv[s],i)%mod*inv[n-i*s]%mod*fas(m-i,n-i*s)%mod;else break;
	reverse(f,f+m+1);k=1;while(k<=m+m)k<<=1;
	NTT(f,1);NTT(g,1);
	for(register int i=0;i<k;i++)f[i]=1ll*f[i]*g[i]%mod;
	NTT(f,-1);INV=fas(k,mod-2);
	for(register int i=0;i<k;i++)f[i]=1ll*f[i]*INV%mod;
	for(register int i=0;i<=m;i++)c[i]=1ll*d[i]*f[i+m]%mod;
	for(register int i=0;i<=m;i++)Add(ans,1ll*c[i]*w[i]%mod);
	printf("%d\n",ans);
	return 0;
}