【UOJ 450】复读机

传送门

---

problem

群里有 $k$ 个不同的复读机。为了庆祝平安夜的到来，在接下来的 $n$ 秒内，它们每秒钟都会选出一位优秀的复读机进行复读。非常滑稽的是，一个复读机只有总共复读了 $d$ 的倍数次才会感到快乐。问有多少种不同的安排方式使得所有的复读机都感到快乐。

答案对 $19491001$ 取模。

数据范围：$n\le 10^9$。
当 $d=1,2$ 时，$k\le 5\times 10^5$。
当 $d=3$ 时，$k\le 1000$。

---

solution

这个数据范围提示我们应该要分类讨论。

---

当 $d=1$ 时，答案显然是 $k^n$。

---

当 $d=2$ 时，我们先把每个复读机的生成函数构造出来：

$$f(x)=\sum _{i=0}^{\infty }\frac{[2|i]}{i!}{x}^i=\frac{e^x+e^{-x}}{2}$$

PS：这个根据 $e^x=\sum _{i=0}^{\infty }\frac{x^i}{i!}$ 就可以推出来。

那么答案应该是 $[x^n]f^k(x)$，我们把 $f^k(x)$ 展开：

$$\begin{array}{rl}{f}^k(x)& =\frac{1}{2^k}(e^x+e^{-x}{)}^k\\ & =\frac{1}{2^l}\sum _{i=0}^k\left(\genfrac{}{}{0ex}{}{k}{i}\right)e^{ix}{e}^{(k-i)(-x)}\\ & =\frac{1}{2^k}\sum _{i=0}^k\left(\genfrac{}{}{0ex}{}{k}{i}\right)e^{(2i-k)x}\end{array}$$

虽然 $[x^n]e^{ax}$ 根据泰勒展开应该是 $\frac{a^n}{n!}$，但是由于 EGF 算出答案后会乘个 $n!$，所以：

$$ans=\frac{1}{2^k}\sum _{i=0}^k\left(\genfrac{}{}{0ex}{}{k}{i}\right)(2i-k{)}^n$$

那么这一部分我们就可以 $O(k\log n)$ 解决掉啦。

---

当 $d=3$ 时。

还是先把生成函数写出来吧（其中 $d=3$）：

$$f(x)=\sum _{i=0}^{\infty }[d|i]\frac{x^i}{i!}$$

看到那个 $[d|i]$ 有点不爽，想办法把它换掉。

考虑单位根反演，即：

$$[d|k]=\frac{{\sum }_{i=0}^{d-1}{\omega }_d^{ik}}{d}$$

简单证明如下：

• 当 $d|k$ 时，根据消去引理，有 ${\omega }_d^{ik}={\omega }_1^{k^{\prime }}=1$，那么 $\frac{{\sum }_{i=0}^{d-1}{\omega }_d^{ik}}{d}=1=$ 左边。
• 当 $d|̸k$ 时，右边相当于是等比数列求和，为 $\frac{1}{d}\times \frac{1-({\omega }_d^k{)}^d}{1-{\omega }_d^k}$，由于 $1-({\omega }_d^k{)}^d=0$，所以右边 $=0=$ 左边。

证毕。

那么代进去，得到：

$$\begin{array}{rl}f(x)& =\frac{1}{d}\sum _{i=0}^{\infty }\sum _{j=0}^{d-1}\frac{({\omega }_d^{j}x{)}^i}{i!}\\ & =\frac{1}{d}\sum _{j=0}^{d-1}{e}^{{\omega }_d^{j}x}\end{array}$$

也就是说，这时的 $f(x)=\frac{1}{3}(e^x+e^{{\omega }_3^{1}x}+e^{{\omega }_3^{2}x})$。

我们把 $f^k(x)$ 暴力展开（就是二项式定理套个二项式定理），得到：

$$f^k(x)=\frac{1}{3^k}\sum _{i=0}^k\left(\genfrac{}{}{0ex}{}{k}{i}\right)\left(\sum _{j=0}^{k-i}\left(\genfrac{}{}{0ex}{}{k-i}{j}\right)e^{ix+{\omega }_3^{1}jx+{\omega }_3^2(k-i-j)x}\right)$$

时间复杂度 $O(k^2)$。

Tip: 处理 ${\omega }_3^1$ 和 ${\omega }_3^2$。

const int P=19491001,g=7;
int w1=power(g,(P-1)/3),w2=mul(w1,w1);

注意要满足 $3|(P-1)$。

---

code

#include<cstdio>
#include<cstring>
#include<algorithm>
#define N 500005
#define P 19491001
using namespace std;
int n,k,d,fac[N],ifac[N];
const int inv2=9745501,inv3=12994001;
const int g=7,w1=18827933,w2=663067;
int add(int x,int y)  {return x+y>=P?x+y-P:x+y;}
int dec(int x,int y)  {return x-y< 0?x-y+P:x-y;}
int mul(int x,int y)  {return 1ll*x*y%P;}
int C(int n,int m)  {return mul(fac[n],mul(ifac[m],ifac[n-m]));}
int power(int a,int b,int ans=1){
	for(;b;b>>=1,a=mul(a,a))
		if(b&1)  ans=mul(ans,a);
	return ans;
}
void init_fac(){
	fac[0]=fac[1]=1;
	for(int i=2;i<N;++i)  fac[i]=mul(fac[i-1],i);
	ifac[N-1]=power(fac[N-1],P-2);
	for(int i=N-2;~i;--i)  ifac[i]=mul(ifac[i+1],i+1);
}
int main(){
	init_fac();
	scanf("%d%d%d",&n,&k,&d);
	if(d==1)  printf("%d\n",power(k,n));
	else  if(d==2){
		int ans=0;
		for(int i=0;i<=k;++i)  ans=add(ans,mul(C(k,i),power(dec(add(i,i),k),n)));
		printf("%d\n",mul(ans,power(inv2,k)));
	}
	else{
		int ans=0;
		for(int i=0;i<=k;++i){
			int temp,sum=0;
			for(int j=0;j<=k-i;++j){
				temp=add(i,add(mul(w1,j),mul(w2,k-i-j)));
				sum=add(sum,mul(C(k-i,j),power(temp,n)));
			}
			ans=add(ans,mul(C(k,i),sum));
		}
		ans=mul(ans,power(inv3,k));
		printf("%d\n",ans);
	}
	return 0;
}