P4491 [HAOI2018]染色 广义容斥 NTT 生成函数

LINK：染色

算是比较常规的广义容斥。

算恰好k个 可以直接转成至少k个。

至少k个非常的好求 直接生成函数。

设\(g_k\)表示至少有k个颜色是满足的 那么有 \(g_k=C(m,k)\frac{n!}{(s!)^k}\frac{(m-k)^{n-sk}}{(n-sk)!}\)

设\(f_k\)表示恰好有k个颜色是满足的 那么有 \(f_k=\sum_{j=k}C(j,k)(-1)^{j-k}g_j\)

前者可以直接求 后者需要卷积一下。

坑点：模数不是998244353 是1004535809 原根也是3. NTT的时候 减法的时候由于数组中有的值可能为负数 所以此时需要也强制转换！

const int MAXN=300010,maxn=10000010,GG=3;
int n,m,s,lim,maxx;
int w[MAXN],rev[MAXN],g[MAXN],f[MAXN];
int inv[maxn],fac[maxn],O[MAXN];
inline int ksm(int b,int p)
{
	int cnt=1;
	while(p)
	{
		if(p&1)cnt=(ll)cnt*b%mod;
		b=(ll)b*b%mod;p=p>>1;
	}
	return cnt;
}
inline int C(int a,int b)
{
	return a<b?0:(ll)fac[a]*inv[b]%mod*inv[a-b]%mod;
}
inline void NTT(int *a,int op)
{
	rep(1,lim-1,i)if(i<rev[i])swap(a[i],a[rev[i]]);
	for(int len=2;len<=lim;len=len<<1)
	{
		int mid=len>>1;
		int wn=ksm(GG,op==1?(mod-1)/len:mod-1-(mod-1)/len);
		rep(1,mid,i)O[i]=(ll)O[i-1]*wn%mod;
		for(int j=0;j<lim;j+=len)
		{
			rep(0,mid-1,i)
			{
				int x=a[i+j],y=(ll)a[i+j+mid]*O[i]%mod;
				a[i+j]=(x+y)%mod;a[i+j+mid]=((ll)x-y+mod)%mod;
			}
		}
	}
	if(op==-1)
	for(int i=0,INV=ksm(lim,mod-2);i<lim;++i)a[i]=(ll)a[i]*INV%mod;
}
signed main()
{
	//freopen("1.in","r",stdin);
	get(n);get(m);get(s);O[0]=fac[0]=1;
	rep(0,m,i)get(w[i]);maxx=max(max(s,m),n);
	rep(1,maxx,i)fac[i]=(ll)fac[i-1]*i%mod;
	inv[maxx]=ksm(fac[maxx],mod-2);
	fep(maxx-1,0,i)inv[i]=(ll)inv[i+1]*(i+1)%mod;
	int w2=1;
	rep(0,m,k)
	{
		if(n<s*k)break;
		g[k]=(ll)C(m,k)*fac[n]%mod*w2%mod;
		g[k]=(ll)g[k]*ksm(m-k,n-s*k)%mod*inv[n-s*k]%mod;
		w2=(ll)w2*inv[s]%mod;
	}
	rep(0,m,i)f[i]=((m-i)&1?-1:1)*inv[m-i],g[i]=(ll)fac[i]*g[i]%mod;
	lim=1;while(lim<=m+m)lim=lim<<1;
	rep(0,lim-1,i)rev[i]=rev[i>>1]>>1|(i&1?lim>>1:0);
	NTT(f,1);NTT(g,1);
	rep(0,lim-1,i)f[i]=(ll)f[i]*g[i]%mod;
	NTT(f,-1);int ans=0;
	rep(0,m,i)ans=(ans+(ll)w[i]*f[i+m]%mod*inv[i])%mod;
	put((ans+mod)%mod);return 0;
}