[HAOI2018]染色

Link

Solution

如果令 (G(i)) 表示恰好有 (i) 种颜色出现了 (S) 次，答案就是

[sum_{i=0}^{min(n/s,m)} w_i imes G(i) ]

令 (lim=min(n/s,m))，我们只需要求出 (0) ~ (lim) 的 (G)。
容易想到令 (F(i)) 表示强制令 (i) 种颜色出现了 (s) 次，剩下的随便选的可重方案，那么

[F(i)=inom{m}{i}inom{n}{si}frac{(si)!}{(s!)^i}(m-i)^{n-si} ]

可以发现，序列 (F) 能够直接 (O(n log n)) 预处理出来。
又容易知道

[G(k)=sum_{i=k}^{lim} (-1)^{i-k} inom{i}{k} F(i)=frac{1}{k!}sum_{i=k}^{lim} frac{(-1)^{i-k}}{(i-k)!} imes i!F(i) ]

再令 (A(i)=frac{(-1)^i}{i!}) 和 (B(i)=(lim-i)!F(lim-i)!)，即 (B(lim-i)=i!F(i))。所以上式可以变成

[G(k)=sum_{i=0}^{lim-k} frac{(-1)^i}{i!} imes (i+k)!F(i+k)=sum_{i=0}^{lim-k} A(i) imes B(lim-i-k) ]

容易发现这是卷积的形式，用 NTT 做一下就行了。求出 (A) 和 (B) 的卷积多项式后，再把其 reverse 一下，就是 (G) 了。

复杂度 (O(nlog n))，其中 (n) 和上文的 (lim) 同阶。

#include<stdio.h>
#include<algorithm>
using namespace std;
#define ll long long
#define rint register int

const int V=10100007;
const int N=(1<<21)+3;
const int P=1004535809;
const int G=3;

inline ll qpow(ll x,ll y){
    ll ret=1,cnt=0;
    while(y>=(1ll<<cnt)){
        if(y&(1ll<<cnt)) ret=ret*x%P;
        x=x*x%P,cnt++;
    }
    return ret;
}

const int Gi=qpow(G,P-2);

inline int read(){
    int x=0,flag=1; char c=getchar();
    while(c<'0'||c>'9'){if(c=='-') flag=0;c=getchar();}
    while(c>='0'&&c<='9'){x=(x<<1)+(x<<3)+c-48;c=getchar();}
    return flag? x:-x;
}

int rk[N],n,m,op,s;
ll a[N],b[N],w[N],fac[V],inv[V];

ll C(int x,int y){return x<y? 0:fac[x]*inv[y]%P*inv[x-y]%P;}

inline void swap(ll &x,ll &y){x^=y,y^=x,x^=y;}
inline void NTT(ll *F){
    for(rint i=0;i<n;i++)
        if(i<rk[i]) swap(F[i],F[rk[i]]);
    for(rint p=2;p<=n;p<<=1){
        int len=p>>1;
        ll w=qpow(op? G:Gi,(P-1)/p);
        for(rint k=0;k<n;k+=p){
            ll now=1;
            for(rint l=k;l<k+len;l++){
                ll t=now*F[l+len]%P;
                F[l+len]=(F[l]-t+P)%P;
                F[l]=(F[l]+t)%P;
                now=now*w%P;
            }
        }
    }
}

inline int min(int x,int y){return x<y? x:y;}
int main(){
    n=read(),m=read(),s=read();
    for(rint i=0;i<=m;i++) w[i]=read();
    int lim=min(m,n/s),M=max(n,m)+1;
    fac[0]=1;
    for(rint i=1;i<M;i++) fac[i]=fac[i-1]*i%P;
    inv[M-1]=qpow(fac[M-1],P-2);
    for(rint i=M-2;~i;i--) inv[i]=inv[i+1]*(i+1)%P;
    for(rint i=0;i<=lim;i++)
        a[lim-i]=fac[i]*C(m,i)%P*fac[n]%P*inv[n-s*i]%P*qpow(inv[s],i)%P*qpow(m-i,n-s*i)%P;
    for(rint i=0,opt=1;i<=lim;i++,opt=-opt) b[i]=(opt*inv[i]+P)%P;
    for(n=1;n<=(lim<<1);n<<=1);
    for(rint i=0;i<n;i++)
        rk[i]=(rk[i>>1]>>1)|((i&1)? n>>1:0);
    op=1,NTT(a),NTT(b);
    for(rint i=0;i<n;i++) a[i]=a[i]*b[i]%P;
    op=0,NTT(a);
    reverse(a,a+1+lim); 
    ll ans=0,Inv=qpow(n,P-2);
    for(rint i=0;i<=lim;i++) ans=(ans+w[i]*a[i]%P*inv[i]%P*Inv%P)%P;
    printf("%lld",ans);
}

Reflection

一开始没想到 (F) 可以直接跑出来，而是很憨的把它展开再化简，什么都没搞出来。以后遇到一元函数一定要先看能不能快速预处理，免得对后面的式子造成麻烦。

最后一步一听说是叫是什么差卷积？不太懂。反正把 (i!F(i)) 反转成 (B) 这一步是真巧妙。似乎二项式反演大概都能这么搞，学到了。