[BZOJ5306] [HAOI2018]染色(容斥原理+NTT)
题面
一个长度为 n的序列, 每个位置都可以被染成 m种颜色中的某一种.
如果n个位置中恰好出现了 S次的颜色有 K种, 则小 C 会产生 (W_k)的愉悦度.
求对于所有可能的染色方案, 他能获得的愉悦度的和。答案对 1004535809 取模
分析
显然颜色数量不超过(tot=min(m,frac{n}{S}))
我们需要求出现了(S)次的颜色有(i)种的方案数。这个东西不太好求，考虑容斥，求出现了(S)次的颜色有至少(i)种的方案数,记为(t(i))
我们先从(m)种颜色里选出(i)种颜色，方案数(C_m^i)
然后从(n)个位置里选出(S)个位置，涂上颜色1。从剩下(n-S)个位置里选出(S)个位置，涂上颜色2。从剩下(n-2S)个位置里选出(S)个位置，涂上颜色3......从从剩下(n-(i-1)S)个位置里选出(S)个位置，涂上颜色i. 方案数为:
[egin{aligned} &= C_n^S cdot C_{n-S}^S cdot C_{n-2S}^S dots C_{n-(i-1)S}^S \ &=frac{n!}{S!(n-S)!} cdot frac{(n-S)!}{S!(n-2S)!} cdot frac{(n-2S)!}{S!(n-3S)!} dots frac{(n-(i-1)S)!}{S!(n-iS)!} \ &=frac{n!}{(S!)^i (n-iS)!} (上式左边的分母可以和右边的分子消掉)end{aligned}
]
这样涂完色之后还有(n-iS)个位置没涂，剩下(m-i)种颜色没用，方案数为((m-i)^{n-iS})
因此，
[t(i)=C_m^i frac{n!}{(S!)^i (n-iS)!} (m-i)^{n-iS}
]
然后加上容斥,最终答案(ans(i))为
[egin{aligned} ans(i)&=sum_{j=i}^{tot} (-1)^{j-i} C_{j}^{i} t(j) \ &=sum_{j=i}^{tot} (-1)^{j-i} frac{j!}{i!(j-i)!}t(j) \ ans(i)i! &=sum_{j=i}^{tot} frac{(-1)^{j-i}}{(j-i)!} cdot t(j)j! end{aligned}
]
令(A(i)=t(i)i!,B(i)=frac{(-1)^{j-i}}{(j-i)!})
[ans(i)i!=sum_{j=i}^{tot} A(j) B(j-i)
]
由于(j+j-i)不是常数，不能直接卷积。类似 [ZJOI2014]力 ，把B反转就能卷积了.
令(B'(j)=B(tot-j)),则(B(j-i)=B'(tot-j+i))
[ans(i)i!=sum_{j=i}^{tot} A(j) B'(tot-j+i)
]
由于(tot-j+i+j=tot+i)为一常数，符合卷积的特征。注意最后的(ans(i)i!)实际上是(A)和(B')卷积得到的多项式第(tot+i)项。最后别忘了除以(i!)并乘上(W_i)
代码
//https://www.luogu.org/problem/P4491
#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#define maxn 10000000
#define maxm 400000
#define G 3
#define mod 1004535809
#define invG 334845270
using namespace std;
typedef long long ll;
inline ll fast_pow(ll x,ll k){
ll ans=1;
while(k){
if(k&1) ans=ans*x%mod;
x=x*x%mod;
k>>=1;
}
return ans;
}
inline ll inv(ll x){
return fast_pow(x,mod-2);
}
void NTT(ll *x,int n,int type){
static int rev[maxm+5];
int tn=1,k=0;
while(tn<n){
tn*=2;
k++;
}
for(int i=0;i<n;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(k-1));
for(int i=0;i<n;i++) if(i<rev[i]) swap(x[i],x[rev[i]]);
for(int len=1;len<n;len*=2){
int sz=len*2;
ll gn1=fast_pow((type==1?G:invG),(mod-1)/sz);
for(int l=0;l<n;l+=sz){
int r=l+len-1;
ll gnk=1;
for(int i=l;i<=r;i++){
ll tmp=x[i+len];
x[i+len]=(x[i]-gnk*tmp%mod+mod)%mod;
x[i]=(x[i]+gnk*tmp%mod)%mod;
gnk=gnk*gn1%mod;
}
}
}
if(type==-1){
ll invn=inv(n);
for(int i=0;i<n;i++) x[i]=x[i]*invn%mod;
}
}
void mul(ll *a,ll *b,ll *ans,int n){
NTT(a,n,1);
NTT(b,n,1);
for(int i=0;i<n;i++) ans[i]=a[i]*b[i]%mod;
NTT(ans,n,-1);
}
ll fact[maxn+5],invfact[maxn+5];
void ini_fact(int n){
fact[0]=1;
for(int i=1;i<=n;i++) fact[i]=fact[i-1]*i%mod;
invfact[n]=inv(fact[n]);
for(int i=n-1;i>=0;i--) invfact[i]=invfact[i+1]*(i+1)%mod;
}
inline ll C(int n,int m){
return fact[n]*invfact[n-m]%mod*invfact[m]%mod;
}
int n,m,s;
ll w[maxm+5];
ll a[maxm+5],b[maxm+5],sum[maxm+5];
int main(){
// printf("%lld
",inv(3));
scanf("%d %d %d",&n,&m,&s);
for(int i=0;i<=m;i++) scanf("%lld",&w[i]);
ini_fact(max(n,m));
int lim=min(n/s,m);
int dn=1;
while(dn<=(lim+1)*2) dn*=2;
for(int i=0;i<=lim;i++){
a[i]=C(m,i)*fact[n]%mod*inv(fast_pow(fact[s],i))%mod*invfact[n-s*i]%mod*fast_pow(m-i,n-s*i)%mod*fact[i]%mod;
//f(x)=C(m,i)*(n!)/((s!)^i(n-si)!)*(m-i)^(n-si)
//A(x)=f(x)x!
}
for(int i=0;i<=lim;i++){
b[i]=(fast_pow(-1,lim-i)*invfact[lim-i]%mod+mod)%mod;
//B(x)=(-1)^x/x!
//类似[ZJOI2014]力,要把B反转
}
mul(a,b,sum,dn);
ll ans=0;
for(int i=0;i<=lim;i++){
ans=(ans+sum[i+lim]*invfact[i]%mod*w[i]%mod)%mod;
}
printf("%lld
",ans);
}