递推序列第n项求解优化：n>1e6及1e12时超时/超内存求助

当前算法逻辑没问题，但当n增大到10^{6及以上（题目要求1<=n<=10}12）时，会超时或者超内存。最初把元素存vector里，后来改成复用变量，过了3个初始测试，但剩下63个全挂了。

递推公式：$A_i = (A_{i-1} + 2 * A_{i-2} + 3 * A_{i-3}) \mod M$，其中$M=10^9+7$。

时间限制：1秒，内存限制：256MB。

当前C++代码如下：

#include<iostream>
#include<cmath>

using namespace std;
using ull = unsigned long long;

ull func(ull n){
    ull a = 1;
    ull b = 1;
    ull c = 2;
    if (n < 2) return a;
    if (n == 3) return c;
    ull res = 0;
    for (ull i = 0; i < n - 3; i++){
        res = (3 * a + 2 * b + c) % (ull)(pow(10, 9) + 7);
        a = b;
        b = c;
        c = res;
    }
    return c;
}

int main() {
    int x; 
    cin >> x;
    cout << func(x);
}

初始测试用例：

• 测试1：输入6，输出34
• 测试2：输入10，输出1096
• 测试3：输入500，输出340736120

请问得换算法还是有其他提速方法？

肯定得换算法，你现在的线性递推是O(n)的时间复杂度，完全处理不了n=1e12的情况——1秒最多跑大概1e8次循环，1e12比这个大了4个数量级，根本不可能跑完。

正确思路：矩阵快速幂（O(logn)时间复杂度） #

对于这种三阶线性递推，可以转成矩阵乘法的形式，然后用快速幂来加速计算：

\begin{bmatrix}
1 & 2 & 3 \
1 & 0 & 0 \
0 & 1 & 0
\end{bmatrix}
\times
\begin{bmatrix}
A_{i-1} \
A_{i-2} \
A_{i-3}
\end{bmatrix}
$$

当n>=3时，计算这个矩阵的(n-3)次幂，再乘以初始向量$\begin{bmatrix}A_3 \ A_2 \ A_1\end{bmatrix} = \begin{bmatrix}2 \ 1 \ 1\end{bmatrix}$，就能得到$A_n$的值。

代码里的坑要先填 #

1. 别用pow(10,9)+7算模数，直接写1000000007就行，pow是浮点数函数，容易有精度误差，导致取模出错。
2. 矩阵乘法和快速幂过程中要随时取模，避免溢出（用unsigned long long能降低溢出风险，但取模不能省）。

完整代码示例 #

#include <iostream>
using namespace std;

typedef unsigned long long ull;
const ull MOD = 1000000007;

// 矩阵乘法：a * b = res
void multiply(ull a[3][3], ull b[3][3], ull res[3][3]) {
    ull temp[3][3] = {0};
    for (int i = 0; i < 3; i++) {
        for (int k = 0; k < 3; k++) {
            if (a[i][k] == 0) continue; // 优化：跳过0元素减少计算
            for (int j = 0; j < 3; j++) {
                temp[i][j] = (temp[i][j] + a[i][k] * b[k][j]) % MOD;
            }
        }
    }
    // 把temp复制到res
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            res[i][j] = temp[i][j];
        }
    }
}

// 矩阵快速幂：mat^power = res
void matrix_pow(ull mat[3][3], ull power, ull res[3][3]) {
    // 初始化res为单位矩阵
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            res[i][j] = (i == j) ? 1 : 0;
        }
    }
    while (power > 0) {
        if (power % 2 == 1) {
            ull temp[3][3];
            multiply(res, mat, temp);
            // 更新res
            for (int i = 0; i < 3; i++) {
                for (int j = 0; j < 3; j++) {
                    res[i][j] = temp[i][j];
                }
            }
        }
        ull temp[3][3];
        multiply(mat, mat, temp);
        // 更新mat为平方后的结果
        for (int i = 0; i < 3; i++) {
            for (int j = 0; j < 3; j++) {
                mat[i][j] = temp[i][j];
            }
        }
        power /= 2;
    }
}

ull func(ull n) {
    if (n == 1 || n == 2) return 1;
    if (n == 3) return 2;
    ull mat[3][3] = {
        {1, 2, 3},
        {1, 0, 0},
        {0, 1, 0}
    };
    ull power = n - 3;
    ull mat_pow[3][3];
    matrix_pow(mat, power, mat_pow);
    // 初始向量是[A3, A2, A1] = [2,1,1]，计算结果的第一行
    return (mat_pow[0][0] * 2 + mat_pow[0][1] * 1 + mat_pow[0][2] * 1) % MOD;
}

int main() {
    ull x;
    cin >> x;
    cout << func(x) << endl;
    return 0;
}

为啥这个方法可行？ #

矩阵快速幂的时间复杂度是O(logn)，log2(1e12)大概是40，也就是说只需要几十次矩阵乘法就能出结果，完全在1秒的时间限制内。而且内存只用固定大小的3x3矩阵，不会随n变大而增加，绝对满足内存要求。

内容的提问来源于stack exchange，提问作者lilof