【组合计数+NTT优化卷积】BZOJ5306 [HAOI2018] 染色

【题目】
lydsy
一个长度为$n$的序列，每个位置可以被染成$m$种颜色中的一种。若一种方案中出现次数恰好为$S$的颜色数有$K$种，则会有$W_K$的愉悦值。问所有方案的愉悦值总和对$1004535809$取模的结果。

$n\le 10^7,m\le 10^5,S\le 150$

【解题思路】
首先这个模数就很$\text{NTT}$了（$1004535809=2^{21}\cdot 479+1$，原根为$3$）

然后我就发现我不会做了。于是以下来自各种$\text{blog}$

设$N=\min (m,\lfloor \frac{n}{S}\rfloor )$表示出现$S$次的颜色数上界，设$f_i$表示染色后出现了$S$次的颜色有$i$种的方案数，则答案就是：
$$\sum _{i=0}^N{w}_i\cdot f_i$$
这个$f_i$直接求比较难求，于是考虑这样一个容斥，设$g_i=i!$：
$$f_i=\left(\genfrac{}{}{0ex}{}{m}{i}\right)\left(\genfrac{}{}{0ex}{}{n}{iS}\right)\frac{(g_{iS})!}{g_S^i}\sum _{j=0}^{N-i}(-1{)}^j(\genfrac{}{}{0ex}{}{m-i}{j}\left)\right(\genfrac{}{}{0ex}{}{n-iS}{jS})\frac{g_{jS}}{g_S^j}(m-i-j{)}^{n-iS-jS}$$
这个柿子的意思大概是先从$m$种颜色中选出$i$种。然后从这$n$个位置中选出$iS$个进行可重排列，就是$\left(\genfrac{}{}{0ex}{}{n}{iS}\right)\frac{(g_{iS})!}{g_S^i}$了。接下来剩下的$m-i$种颜色，在$n-iS$个位置上随便填，但由于可能会使某种颜色出现了$S$次，所以后面这个东西我们需要容斥。容斥的意义就是额外出现了$S$次的颜色有$j$种，我们就需要在$m-i$种颜色中选出这$j-i$种，并在$n-iS$个位置中选出$jS$个进行可重排列，剩下的位置随便填。

这样$(m-i-j{)}^{n-iS-jS}$不太好处理，不妨将第二个和式中的$j$意义改变为实际上有$j$种颜色出现了$S$次，那么就是用$j-i$替换$j$并改变循环下标：
$$f_i=\left(\genfrac{}{}{0ex}{}{m}{i}\right)\left(\genfrac{}{}{0ex}{}{n}{iS}\right)\frac{(g_{iS})!}{g_S^i}\sum _{j=i}^N(-1{)}^{j-i}(\genfrac{}{}{0ex}{}{m-i}{j-i}\left)\right(\genfrac{}{}{0ex}{}{n-iS}{jS-iS})\frac{g_{jS-iS}}{g_S^{j-i}}(m-j{)}^{n-jS}$$
接下来就是化简这个柿子，那么将组合数写成阶乘形式后化简：
$$f_i=\frac{g_m\cdot g_n}{g_i}\sum _{j=i}^N\frac{(-1{)}^{j-i}\cdot (m-j{)}^{n-jS}}{g_{j-i}\cdot g_{m-j}\cdot g_{n-jS}\cdot g_S^j}$$
这里面的$j$不太好做，不妨把$j$提到外层，那么最后的答案可以得到：
$$\begin{array}{rl}ans& =\sum _{i=0}^N{w}_i\cdot f_i\\ & =\sum _{i=0}^N{w}_i\cdot \frac{g_m\cdot g_n}{g_i}\sum _{j=i}^N\frac{(-1{)}^{j-i}\cdot (m-j{)}^{n-jS}}{g_{j-i}\cdot g_{m-j}\cdot g_{n-jS}\cdot g_S^j}\\ & =g_m\cdot g_n\cdot \sum _{j=0}^N\frac{(m-j{)}^{n-jS}}{g_{m-j}\cdot g_{n-jS}\cdot g_S^j}\sum _{i=0}^j\frac{(-1{)}^{j-i}\cdot w_i}{g_i\cdot g_{j-i}}\end{array}$$
现在左边是一个只与$j$有关的柿子，右边是$A(x)=\frac{w_x}{g_x}$和$B(x)=\frac{(-1{)}^x}{g_x}$的卷积。

那么我们就可以愉快地做$\text{NTT}$了，最后复杂度$O(n+m\log m)$

【参考代码】荣获BZOJ垫底

#include<bits/stdc++.h>
using namespace std;

typedef long long ll;
const int N=1e7+10,M=262333,mod=1004535809,G=3;

namespace IO
{
	int read()
	{
		int ret=0;char c=getchar();
		while(!isdigit(c)) c=getchar();
		while(isdigit(c)) ret=ret*10+(c^48),c=getchar();
		return ret;
	}
	void write(int x){if(x>9)write(x/10);putchar(x%10^48);}
	void writeln(int x){write(x);putchar('\n');}
}
using namespace IO;

namespace Math
{
	int fac[N],ifac[N];
	int upm(int x){return x>=mod?x-mod:(x<0?x+mod:x);}
	void up(int &x,int y){x=upm(x+y);}
	int mul(int x,int y){return 1ll*x*y%mod;}
	int qpow(int x,int y)
	{
		int res=1;
		for(;y;y>>=1,x=mul(x,x)) if(y&1) res=mul(res,x);
		return res;
	}
	int C(int x,int y){if(y>x)return 0;return mul(fac[x],mul(ifac[x-y],ifac[y]));}
	void initsth()
	{
		fac[0]=1;for(int i=1;i<N;++i)fac[i]=mul(fac[i-1],i);
		ifac[N-1]=qpow(fac[N-1],mod-2);for(int i=N-2;~i;--i)ifac[i]=mul(ifac[i+1],i+1);
	}
}
using namespace Math;

namespace Poly
{
	int m,L;
	int rev[M];
	void ntt(int *a,int n,int op)
	{
		for(int i=0;i<n;++i) if(i<rev[i]) swap(a[i],a[rev[i]]);
		for(int i=1;i<n;i<<=1)
		{
			int wn=qpow(G,(mod-1)/(i<<1));
			if(!~op) wn=qpow(wn,mod-2);
			for(int j=0;j<n;j+=(i<<1))
			{
				int w=1;
				for(int k=0;k<i;++k,w=mul(w,wn))
				{
					int x=a[j+k],y=mul(w,a[i+j+k]);
					a[j+k]=upm(x+y);a[i+j+k]=upm(x-y);
				}
			}
		}
		if(!~op) for(int inv=qpow(n,mod-2),i=0;i<n;++i) a[i]=mul(a[i],inv);
	}
	void mulpoly(int *a,int *b,int *c,int n)
	{
		for(m=1,L=0;m<=n;m<<=1,++L);
		for(int i=0;i<m;++i) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
		ntt(a,m,1);ntt(b,m,1);
		for(int i=0;i<m;++i) c[i]=mul(a[i],b[i]);
		ntt(c,m,-1);
	}
}
using namespace Poly;

namespace DreamLolita
{
	int n,m,S,lim;
	int w[M],f[M],g[M];
	void input()
	{
		n=read();m=read();S=read();lim=min(m,n/S);
		for(int i=0;i<=m;++i) w[i]=read();
	}
	void initf()
	{	
	 	for(int i=0;i<=lim;++i) f[i]=mul(w[i],ifac[i]),g[i]=upm((i&1?-1:1)*ifac[i]);
	 	mulpoly(f,g,f,lim*2);
	}
	void getans()
	{
		int ans=0;
		for(int j=0;j<=lim;++j)
		{
			int t1=qpow(m-j,n-j*S),t2=mul(ifac[m-j],mul(ifac[n-j*S],qpow(ifac[S],j)));
			up(ans,mul(mul(t1,t2),f[j]));
		}
		ans=mul(ans,mul(fac[m],fac[n]));
		writeln(ans);
	}
	void solution(){initsth();input();initf();getans();}
}

int main()
{
#ifdef Durant_Lee
	freopen("BZOJ5306.in","r",stdin);
	freopen("BZOJ5306.out","w",stdout);
#endif
	DreamLolita::solution();
	return 0;
}