【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;
}