【UOJ450】【集训队作业2018】复读机（生成函数+单位根反演）

点此看题面

• 求有多少长度为\(n\)、值域为\([1,k]\)的整数数列，满足所有数的出现次数都是\(d\)的倍数。
• \(n\le10^9\)，\(k\le5\times10^5,d\le2\)或\(k\le10^3,d\le3\)

指数型生成函数

考虑用一个指数型生成函数\(F(x)\)来表示每个数填\(i\)个的方案数，那么最终答案就应该是：

\[[\frac{x^n}{n!}]F(x)^k \]

而这个生成函数其实是非常显然的，就是\(\sum\frac{x^i}{i!}\)只保留\(d\)的倍数项：

\[F(x)=\sum_{i=0}^{+\infty}[d|i]\frac{x^i}{i!} \]

单位根反演

考虑单位根反演的式子：

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

代入就可以得到：

\[F(x)=\frac1d\sum_{j=0}^{d-1}\sum_{i=0}^{+\infty}\frac{(\omega_d^jx)^i}{i!} \]

发现后面的\(\sum\)中的一项可以直接看作\(e\)的幂次，因此得到：

\[F(x)=\frac1d\sum_{j=0}^{d-1}e^{\omega_d^jx} \]

暴力枚举

要求\(F(x)^k\)，先提出每一项的\(\frac 1d\)共计\((\frac1d)^k\)，而剩余部分我们只需要考虑每次选的\(e^{\omega_d^jx}\)的\(j\)是多少。

不妨直接枚举选择了\(i_j\)个\(e^{\omega_d^j}x\)，那么最终答案就应该是：

\[[\frac{x^n}{n!}](\sum_{\sum_{j=0}^{d-1}i_j=k}\frac{k!}{\prod_{j=0}^{d-1}i_j!}e^{(\sum_{j=0}^{d-1}i_j\times \omega_d^j)x}) \]

前面是一堆系数，其实只用考虑\(e^{(\sum_{j=0}^{d-1}i_j\times \omega_d^j)x}\)的\(n\)次项系数然后乘上\(n!\)，于是得到：

\[\sum_{\sum_{j=0}^{d-1}i_j=k}\frac{k!}{\prod_{j=0}^{d-1}i_j!}(\sum_{j=0}^{d-1}i_j\times \omega_d^j)^n \]

所以说具体实现只要暴力枚举所有\(i_j\)即可。

代码：\(O(k^{d-1}logn)\)

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define K 500000
#define X 19491001
#define PR 7//预先算出19491001的原根是7
using namespace std;
int n,k,d;I int QP(RI x,RI y) {RI t=1;W(y) y&1&&(t=1LL*t*x%X),x=1LL*x*x%X,y>>=1;return t;}
int Fac[K+5],IFac[K+5];I void InitFac(CI S)//预处理阶乘和阶乘逆元
{
	RI i;for(Fac[0]=i=1;i<=S;++i) Fac[i]=1LL*Fac[i-1]*i%X;
	for(IFac[i=S]=QP(Fac[S],X-2);i;--i) IFac[i-1]=1LL*IFac[i]*i%X;
}
int main()
{
	if(scanf("%d%d%d",&n,&k,&d),d==1) return printf("%d\n",QP(k,n)),0;//d=1，答案其实就是k^n
	RI i,j,t=0;if(InitFac(k),d==2)//d=2
	{
		for(i=0;i<=k;++i) t=(t+1LL*Fac[k]*IFac[i]%X*IFac[k-i]%X*QP((2*i-k+X)%X,n))%X;//枚举有i个1，k-1个-1
		return printf("%d\n",1LL*t*QP(QP(d,X-2),k)%X),0;//注意乘上(1/d)^k
	}
	RI w=QP(PR,(X-1)/3),w2=1LL*w*w%X;for(i=0;i<=k;++i) for(j=0;j<=k-i;++j)//枚举有i个w，j个w2，k-i-j个1
		t=(t+1LL*Fac[k]*IFac[i]%X*IFac[j]%X*IFac[k-i-j]%X*QP((1LL*i*w+1LL*j*w2+(k-i-j))%X,n))%X;
	return printf("%d\n",1LL*t*QP(QP(d,X-2),k)%X),0;//注意乘上(1/d)^k
}