【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+ichoose 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{kchoose 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}^ksum_{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^dequiv 1(mod p))。所以有(omega_d=g^{frac{p-1}{d}})。
那么算出来之后就可以(O(k^2log 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
",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
",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
",ans);
	}
	return 0;
}