【BZOJ3328】PYXFIB（单位根反演，矩阵快速幂）

【BZOJ3328】PYXFIB（单位根反演，矩阵快速幂）

题面

BZOJ

题解

首先要求的式子是：\(\displaystyle \sum_{i=0}^n [k|i]{n\choose i}f_i\)。

斐波那契数列如果要快速算显然就只能对应着一个矩阵，所以我们就直接默认\(f_i\)是一个矩阵的形式。

如果没有\([k|i]\)这个东西这个玩意看着就很像一个二项式定义的展开。

实际上二项式定理对于矩阵而言显然是对的，不妨设斐波那契数列的转移矩阵是\(A\)，

那么\(\displaystyle (A+I)^n=\sum_{i=0}^n {n\choose i}A^i\)，其中\(I\)是单位矩阵。

因为有前面那个东西，所以用单位根反演直接拆：

\[\begin{aligned}
\sum_{i=0}^n[k|i]{n\choose i}f_i&=\sum_{i=0}^n{n\choose i}f_i\frac{1}{k}\sum_{j=0}^{k-1}\omega_k^{ji}\\
&=\frac{1}{k}\sum_{j=0}^{k-1}\sum_{i=0}^n{n\choose i}f_i(\omega^j)^i
\end{aligned}\]

似乎把\(\omega\)一起丢进矩阵里面就可以直接计算了。

（也就是后面那个东西是\((A\omega^j+I)^n\)）

#include<iostream>
#include<cstdio>
using namespace std;
#define ll long long
inline ll read()
{
	ll x=0;bool t=false;char ch=getchar();
	while((ch<'0'||ch>'9')&&ch!='-')ch=getchar();
	if(ch=='-')t=true,ch=getchar();
	while(ch<='9'&&ch>='0')x=x*10+ch-48,ch=getchar();
	return t?-x:x;
}
ll n;int K,MOD,g,w,ans;
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;}
struct Matrix
{
	int s[2][2];
	void clear(){s[0][0]=s[0][1]=s[1][0]=s[1][1]=0;}
	void init(){s[0][0]=s[1][1]=1;s[0][1]=s[1][0]=0;}
	int*operator[](int x){return s[x];}
}A,I;
Matrix operator*(Matrix a,Matrix b)
{
	Matrix c;c.clear();
	for(int i=0;i<2;++i)
		for(int j=0;j<2;++j)
			for(int k=0;k<2;++k)
				c[i][j]=(c[i][j]+1ll*a[i][k]*b[k][j])%MOD;
	return c;
}
Matrix operator+(Matrix a,Matrix b)
{
	for(int i=0;i<2;++i)
		for(int j=0;j<2;++j)
			a[i][j]=(a[i][j]+b[i][j])%MOD;
	return a;
}
Matrix operator*(Matrix a,int b)
{
	for(int i=0;i<2;++i)
		for(int j=0;j<2;++j)
			a[i][j]=1ll*a[i][j]*b%MOD;
	return a;
}
Matrix fpow(Matrix a,ll b)
{
	Matrix s;s.init();
	while(b){if(b&1)s=s*a;a=a*a;b>>=1;}
	return s;
}
int fac[100],tot;
int Getg()
{
	int x=MOD-1;tot=0;
	for(int i=2;i*i<=x;++i)
		if(x%i==0)
		{
			fac[++tot]=i;
			while(x%i==0)x/=i;
		}
	if(x>1)fac[++tot]=x;
	for(int i=2;;++i)
	{
		bool fl=true;
		for(int j=1;j<=tot;++j)
			if(fpow(i,(MOD-1)/fac[j])==1){fl=false;break;}
		if(fl)return i;
	}
}
int main()
{
	int T=read();I.init();
	while(T--)
	{
		n=read();K=read();MOD=read();
		g=Getg();w=fpow(g,(MOD-1)/K);ans=0;
		for(int i=0,tw=1;i<K;++i,tw=1ll*tw*w%MOD)
		{
			A[0][0]=1;A[0][1]=1;A[1][0]=1;A[1][1]=0;
			A=fpow(A*tw+I,n);
			ans=(0ll+ans+A[0][1]+A[1][1])%MOD;
		}
		ans=1ll*ans*fpow(K,MOD-2)%MOD;
		printf("%d\n",ans);
	}
	return 0;
}