[HAOI2018]染色

Link

Solution

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

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

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

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

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

\[G(k)=\sum_{i=k}^{lim} (-1)^{i-k} \binom{i}{k} F(i)=\frac{1}{k!}\sum_{i=k}^{lim} \frac{(-1)^{i-k}}{(i-k)!} \times 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!} \times (i+k)!F(i+k)=\sum_{i=0}^{lim-k} A(i)\times B(lim-i-k) \]

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

复杂度 \(O(n\log 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\) 这一步是真巧妙。似乎二项式反演大概都能这么搞，学到了。