洛谷P4491 [HAOI2018]染色

题面
题解：首先这个“恰好”看着很不爽，换成“至少”。
设\(f[i]\)表示颜色个数为\(S\)的颜色至少有\(i\)个的方案数。
考虑如何计算。
1.\(m\)个颜色选了\(i\)个，\(\binom{m}{i}\)；
2.\(i\)个颜色选了恰好\(S\)个，其他的\(m-i\)个颜色任选;\[\frac {n!}{ {S!}^i \times fac[n-i*S]} \]
3.剩下\(n-i*S\)个位置任选\(m-i\)中颜色，\({m-i}^{n-i*S}\)。
把它们乘起来就是\(f[i]\)。
接下来考虑计算答案。
考虑容斥。设\(g[i]\)为颜色个数为\(S\)的颜色恰好有\(i\)个的方案数。于是有：
\[g[i]= \sum_{j=i}^{lim} \frac { {-1}^{j-i} } {fac[j-i]} \times f[j] \times \binom{j}{i}\]
这个\(lim\)表示的是最多能取到的恰好有\(S\)个的颜色数，\(lim=min(n/S,m)\)。
将组合数拆开，试着把式子化为卷积的形式。
\[g[i] \times fac[i] = \sum_{j=i}^{lim} \frac { {-1}^{j-i} } {fac[j-i]} \times f[j] \times fac[j]\]
这样用NTT做就好了。
令\[a[i]=f[i] \times fac[i],b[i]= \frac { {-1}^{j-i} } {fac[j-i]} \]
把\(b\)数组翻转一下，取做完卷积的数组的\(lim\)到\(2lim\)位更新答案即可。
时间复杂度：\(O(n+mlogm)\)
代码：

#include
using namespace std;
#define re register int
#define F(x,y,z) for(re x=y;x<=z;x++)
#define FOR(x,y,z) for(re x=y;x>=z;x--)
typedef long long ll;
#define I inline void
#define IN inline int
#define Cl(x,y) memset(x,y,sizeof(x))
#define kTk system("pause")
templateI read(D &res){
    res=0;register D g=1;register char ch=getchar();
    while(!isdigit(ch)){
        if(ch=='-')g=-1;
        ch=getchar();
    }
    while(isdigit(ch)){
        res=(res<<3)+(res<<1)+(ch^48);
        ch=getchar();
    }
    res*=g;
}
const int Mod=1004535809;
int n,m,k,S,len,lim,ans,r[303000],w[303000],fac[10100000],inv[10100000],f[303000],A[303000],B[303000];
I Add(int &x,int y){
    x+=y;if(x>=Mod)x-=Mod;
}
IN Plus(int x,int y){
    x+=y;if(x>=Mod)x-=Mod;return x;
}
IN Pow(int x,int y){
    re res=1;
    while(y){
        if(y&1)res=(ll)res*x%Mod;
        x=(ll)x*x%Mod;
        y>>=1;
    }
    return res;
}
IN C(int x,int y){
    return (ll)fac[x]*inv[y]%Mod*inv[x-y]%Mod;
}
I init(){
    re M=10001000;
    fac[0]=1;
    F(i,1,M)fac[i]=(ll)fac[i-1]*i%Mod;
    inv[M]=Pow(fac[M],Mod-2);
    FOR(i,M-1,0)inv[i]=(ll)inv[i+1]*(i+1)%Mod;
}
I ntt(int *a,int sn){
    F(i,1,S)if(i>1]>>1)|((i&1)<<(len-1));
    }
    ntt(A,1);ntt(B,1);
    F(i,0,S-1)A[i]=(ll)A[i]*B[i]%Mod;
    ntt(A,-1);
    F(i,lim,lim<<1)Add(ans,(ll)A[i]*w[i-lim]%Mod*inv[i-lim]%Mod);
    printf("%d",ans);
    return 0;
}