【UOJ#450】【集训队作业2018】复读机（生成函数，单位根反演）

【UOJ#450】【集训队作业2018】复读机（生成函数，单位根反演）

题面

UOJ

题解

似乎是\(\mbox{Anson}\)爷的题。
\(d=1\)的时候，随便怎么都行，答案就是\(k^n\)。
\(d=2\)的时候，可以做一个\(dp\)，设\(f[i][j]\)表示前\(i\)个复读机选了\(j\)个时间的方案数。
然后枚举当前这个复读机复读的次数，得到：

\[f[x][j]=\sum_{i=0}^{j}[2|i]{n-j+i\choose i}f[x-1][j-i] \]

化简啥的之后得到

\[(n-j)!f[x][j]=\sum_{i=0}^{j}[2|i]\frac{1}{i!}(n-j+i)!f[x-1][j-i] \]

那么对于\((n-j)!f_j\)构建生成函数，那么等价于每加入一个数就乘上了一个\(A(x)=\sum_{i=0}^{\infty}[2|i]\frac{1}{i!}x^i\)
化简之后就是\(A(x)=\frac{e^x+e^{-x}}{2}\)。
所以最终的答案就是\(A(x)^k[n]\)，即\((\frac{e^x+e^{-x}}{2})^k[x^n]\)。
把那个除二搞出来，拿二项式定理算算得到\(\displaystyle \sum_{i=0}^k{k\choose i}e^{(2i-k)x}\) 。
所以\(n\)次项系数就是\((2i-k)^n\)。
最后就是对于\(d=3\)的情况，
推出来的式子就是

\[\sum_{i=0}^k[3|i]\frac{1}{i!}x^i \]

实际上没必要拿\(dp\)方程来推，可以直接用生成函数考虑。如果确定了每个复读机复读的次数，那么总方案实际上就是\(n!\)除上每个复读机复读次数的阶乘。这个可重排列可以直接推出这个生成函数。

然后发现不会算，前面那个\([3|i]\)不会搞，这样子就可以丢一脸的单位根反演出来：

\[[d|i]=\frac{1}{d}\sum_{j=0}^{d-1}\omega_d^{ij} \]

然后原式就变成了

\[\frac{1}{d}\sum_{i=0}^k\sum_{j=0}^{d-1}\omega_d^{ij}\frac{1}{i!}x^i \]

也就是

\[\frac{1}{d}\sum_{j=0}^{d-1}e^{\omega_d^{j}x} \]

要求的是这个东西的\(k\)次方，也就是

\[\frac{1}{d^k}(\sum_{i=0}^{d-1}e^{\omega_d^{j}x})^k \]

而\(d=3\)的时候\(k\)很小，直接\(O(k^2)\)暴力二项式定理给他展开就好了。
至于单位根怎么求？
求出模数\(p\)的原根\(g\)，我们知道\(g^{p-1}\equiv 1(mod\ p)\)，而\(\omega_d^d\equiv 1(mod\ p)\)。所以有\(\omega_d=g^{\frac{p-1}{d}}\)。
那么算出来之后就可以\(O(k^2\log n)\)计算答案了。

#include<iostream>
#include<cstdio>
using namespace std;
#define MOD 19491001
#define MAX 500500
int fpow(int a,int b)
{
	int s=1;
	while(b){if(b&1)s=1ll*s*a%MOD;a=1ll*a*a%MOD;b>>=1;}
	return s;
}
int n,k,d,ans;
int jc[MAX],jv[MAX],inv[MAX];
int C(int n,int m){if(n<m||n<0)return 0;return 1ll*jc[n]*jv[m]%MOD*jv[n-m]%MOD;}
int main()
{
	scanf("%d%d%d",&n,&k,&d);
	if(d==1){printf("%d\n",fpow(k,n));return 0;}
	jc[0]=jv[0]=inv[0]=inv[1]=1;
	for(int i=2;i<=k;++i)inv[i]=1ll*inv[MOD%i]*(MOD-MOD/i)%MOD;
	for(int i=1;i<=k;++i)jc[i]=1ll*jc[i-1]*i%MOD;
	for(int i=1;i<=k;++i)jv[i]=1ll*jv[i-1]*inv[i]%MOD;
	if(d==2)
	{
		for(int i=0;i<=k;++i)ans=(ans+1ll*C(k,i)*fpow((i+i-k+MOD)%MOD,n))%MOD;
		ans=1ll*ans*fpow(fpow(2,k),MOD-2)%MOD;
		printf("%d\n",ans);
	}
	else
	{
		int w1=fpow(7,(MOD-1)/3),w2=1ll*w1*w1%MOD;
		for(int i=0;i<=k;++i)
			for(int j=0;i+j<=k;++j)
			{
				int p=(1ll*(k-j-i)+1ll*w1*i+1ll*w2*j)%MOD;
				ans=(ans+1ll*C(k,i)*C(k-i,j)%MOD*fpow(p,n))%MOD;
			}
		ans=1ll*ans*fpow(fpow(d,k),MOD-2)%MOD;
		printf("%d\n",ans);
	}
	return 0;
}