【题解】Luogu P5175 数列

题目大意

给定一个递推式 \(a_n=x \times a_{n-1}+ y \times a_{n-2}(n≥3)\)，求 \(\sum_{i=1}^na_i^2\)。

解题思路

递推通常是 \(O(n)\) 解法，但是本题 \(1 \le n \le 10^{18}\) 且 \(T=30000\)（注意是等于），所以一定会超时。这里就要用到矩阵快速幂进行优化。

定义 \(f(n)\) 表示 \(a_n\)，\(s(n)\) 表示 \(\sum_{i=1}^na_i^2\)，则可以根据完全平方公式写出递推式：

\(f(n)=xf(n-1)+yf(n-2)\)

\(f(n)^2=(xf(n-1)+yf(n-2))=x^2f(n-1)^2+y^2f(n-2)^2+2xyf(n-1)f(n-2)\)

\(f(n)f(n-1)=xf(n-1)^2+yf(n-1)f(n-2)\)

\(s(n)=s(n-1)+f(n)^2\)

接下来我们构造状态矩阵。我们发现求 \(s(n)\) 需要用到四个值：\(s(n-1)\)、\(f(n-1)^2\)、\(f(n-2)^2\)、\(f(n-1)f(n-2)\)，所以我们构造两个状态矩阵：

\[\begin{bmatrix} s(n) & f(n)^2 & f(n-1)^2 & f(n)f(n-1) \end{bmatrix} = \begin{bmatrix} s(n-1) & f(n-1)^2 & f(n-2)^2 & f(n-1)f(n-2) \end{bmatrix} \times A \]

而这个 \(A\) 矩阵就是我们要求的转移矩阵。要想求 \(A\) 矩阵，我们先将第一个矩阵的每一项带入刚才求得的递推式来求用第二个矩阵的全部项求第一个矩阵的那一项所需要乘的系数，也就是如何转移：

第一项

\(s(n)=s(n-1) \times 1 + f(n-1)^2 \times x^2 + f(n-2)^2 \times y^2 + f(n-1)f(n-2) \times 2xy\)

第二项

\(f(n)^2=s(n-1) \times 0 + f(n-1)^2 \times x^2 + f(n-2)^2 \times y^2 + f(n-1)f(n-2) \times 2xy\)

第三项

\(f(n-1)^2=s(n-1) \times 0 + f(n-1)^2 \times 1 + f(n-2)^2 \times 0 + f(n-1)f(n-2) \times 0\)

第四项

\(f(n)f(n-1)=s(n-1) \times 0 + f(n-1)^2 \times x + f(n-2)^2 \times 0 + f(n-1)f(n-2) \times y\)

那么我们把这些系数按矩阵乘法的运算方式一列一列填进矩阵 \(A\) 里，就求出了 \(A\)。

\[A=\begin{bmatrix} 1 & 0 & 0 & 0 \\ x^2 & x^2 & 1 & x \\ y^2 & y^2 & 0 & 0 \\ 2xy & 2xy & 0 & y \end{bmatrix} \]

因为 \(n=1\) 和 \(n=2\) 时的矩阵是需要单独求的，所以最终的式子是：

\[\begin{bmatrix} s(n) & f(n)^2 & f(n-1)^2 & f(n)f(n-1) \end{bmatrix} = \begin{bmatrix} s(2) & f(2)^2 & f(1)^2 & f(2)f(1) \end{bmatrix} \times \begin{bmatrix} 1 & 0 & 0 & 0 \\ x^2 & x^2 & 1 & x \\ y^2 & y^2 & 0 & 0 \\ 2xy & 2xy & 0 & y \end{bmatrix}^{n-2} \]

然后套矩阵快速幂的板子即可。

代码实现

#include<iostream>
#include<cstdio>
#include<cstring>
#define int long long //这里图方便就这样写了。平时不建议这样写。
using namespace std;
const int p=1e9+7; //要取模的数。
struct Matrix{ //矩阵。
	int mx[5][5];
}a,s,c;
int T,n,k,f1,f2,x,y;
Matrix operator *(const Matrix &a,const Matrix &b){ //乘法运算符重载为矩阵乘法。注意它只会对指定的结构体生效。
	memset(c.mx,0,sizeof(c.mx));
	for(int i=1;i<=4;i++){
		for(int j=1;j<=4;j++){
			for(int d=1;d<=4;d++){
				c.mx[i][j]=(c.mx[i][j]%p+(a.mx[i][d]*b.mx[d][j])%p)%p;
			}
		}
	}
	return c;
}
signed main(){
	cin>>T;
	while(T--){
		memset(s.mx,0,sizeof(s.mx)); //清空矩阵。
		memset(a.mx,0,sizeof(a.mx));
		scanf("%lld%lld%lld%lld%lld",&n,&f1,&f2,&x,&y);
		f1%=p;f2%=p;x%=p;y%=p;
		for(int i=1;i<=4;i++) s.mx[i][i]=1; //初始化矩阵。s是单位矩阵，a是转移矩阵。
		a.mx[1][1]=a.mx[2][3]=1;
		a.mx[2][1]=a.mx[2][2]=x*x%p;
		a.mx[3][1]=a.mx[3][2]=y*y%p;
		a.mx[4][1]=a.mx[4][2]=2*x*y%p;
		a.mx[2][4]=x;a.mx[4][4]=y;
		if(n==1) printf("%lld\n",(f1*f1)%p); //特殊情况的判定。
		else if(n==2) printf("%lld\n",(f1*f1%p+f2*f2%p)%p);
		else{
			int s2=(f1*f1%p+f2*f2%p)%p; //求出n=2时的矩阵。
			int f22=(f2*f2)%p;
			int f11=(f1*f1)%p;
			int f12=(f1*f2)%p;
			n-=2;
			while(n){
				if(n&1) s=s*a;
				a=a*a;
				n>>=1;
			}
			printf("%lld\n",(((s2*s.mx[1][1])%p+(f22*s.mx[2][1]))%p+((f11*s.mx[3][1])%p+(f12*s.mx[4][1])%p)%p)%p);
		}
	}
	return 0;
}

时间复杂度为 \(O(T\log n \times m^3)\)，\(m\) 为矩阵边长也就是 \(4\)。

注意事项

• 记得开 long long。
• 记得随时取模。
• 记得不要取模 \(n\)（我在这里卡了很长时间，绷）。
• 记得输出一般情况时乘上 \(n=2\) 时的状态矩阵。
• 常数时间复杂度较高，需要卡常。比如：如果用 memset 清空的话矩阵不要开太大、使用 scanf 和 printf 进行输入输出（或快读）等。