一文详解如何使用「矩阵快速幂」优化 DP 过程（含模板 & 完整数学推导）讲完了线性 DP 专题，我们来看看如何使用求解

本文正在参加「金石计划 . 瓜分6万现金大奖」

矩阵快速幂

讲完了线性 DP 专题，我们来看看如何使用求解数据范围在 $1e9$ 甚至以上的线性 DP 问题。

对于常规的线性 DP 题而言，我们往往需至少 $O(n)$ 的复杂度求解，而利用递归式能够归结为矩阵乘法，同时乘法具有结合律，我们往往能够将复杂度优化到 $O(\log{n})$ ，当然「矩阵快速幂」显然属于面试/笔试中的进阶内容。

今天将使用 $4$ 道题目带大家体验「矩阵快速幂」的神奇作用。

1137. 第 N 个泰波那契数

泰波那契序列 Tn 定义如下：

T0 = 0, T1 = 1, T2 = 1, 且在 n >= 0 的条件下 Tn+3 = Tn + Tn+1 + Tn+2

给你整数 $n$ ，请返回第 $n$ 个泰波那契数 $T_n$ 的值。

示例 1：

输入：n = 4

输出：4

解释：
T_3 = 0 + 1 + 1 = 2
T_4 = 1 + 1 + 2 = 4

提示：

0 <= n <= 37
答案保证是一个 32 位整数，即 answer <= $2^{31}$ - 1。

迭代实现动态规划

都直接给出状态转移方程了，其实就是道模拟题。

使用三个变量，从前往后算一遍即可。

代码：

class Solution {
    public int tribonacci(int n) {
        if (n == 0) return 0;
        if (n == 1 || n == 2) return 1;
        int a = 0, b = 1, c = 1;
        for (int i = 3; i <= n; i++) {
            int d = a + b + c;
            a = b;
            b = c;
            c = d;
        }
        return c;
    }
}

时间复杂度： $O(n)$
空间复杂度： $O(1)$

递归实现动态规划

也就是记忆化搜索，创建一个 cache 数组用于防止重复计算。

代码：

class Solution {
    int[] cache = new int[40];
    public int tribonacci(int n) {
        if (n == 0) return 0;
        if (n == 1 || n == 2) return 1;
        if (cache[n] != 0) return cache[n];
        cache[n] = tribonacci(n - 1) + tribonacci(n - 2) + tribonacci(n - 3); 
        return cache[n];
    }
}

时间复杂度： $O(n)$
空间复杂度： $O(n)$

矩阵快速幂

这还是一道「矩阵快速幂」的板子题。

首先你要对「快速幂」和「矩阵乘法」概念有所了解。

矩阵快速幂用于求解一般性问题：给定大小为 $n * n$ 的矩阵 $M$ ，求答案矩阵 $M^k$ ，并对答案矩阵中的每位元素对 $P$ 取模。

在上述两种解法中，当我们要求解 $f[i]$ 时，需要将 $f[0]$ 到 $f[n - 1]$ 都算一遍，因此需要线性的复杂度。

对于此类的「数列递推」问题，我们可以使用「矩阵快速幂」来进行加速（比如要递归一个长度为 $1e9$ 的数列，线性复杂度会被卡）。

使用矩阵快速幂，我们只需要 $O(\log{n})$ 的复杂度。

根据题目的递推关系（ $i >= 3$ ）：

f(i) = f(i - 1) + f(i - 2) + f(i - 3)

我们发现要求解 $f(i)$ ，其依赖的是 $f(i - 1)$ 、 $f(i - 2)$ 和 $f(i - 3)$ 。

我们可以将其存成一个列向量：

\begin{bmatrix} f(i - 1)\\ f(i - 2)\\ f(i - 3) \end{bmatrix}

当我们整理出依赖的列向量之后，不难发现，我们想求的 $f(i)$ 所在的列向量是这样的：

\begin{bmatrix} f(i)\\ f(i - 1)\\ f(i - 2) \end{bmatrix}

利用题目给定的依赖关系，对目标矩阵元素进行展开：

\begin{bmatrix} f(i)\\ f(i - 1)\\ f(i - 2) \end{bmatrix} = \begin{bmatrix} f(i - 1) * 1 + f(i - 2) * 1 + f(i - 3) * 1\\ f(i - 1) * 1 + f(i - 2) * 0 + f(i - 3) * 0\\ f(i - 1) * 0 + f(i - 2) * 1 + f(i - 3) * 0 \end{bmatrix}

那么根据矩阵乘法，即有：

\begin{bmatrix} f(i)\\ f(i - 1)\\ f(i - 2) \end{bmatrix} = \begin{bmatrix} 1 &1 &1 \\ 1 &0 &0 \\ 0 &1 &0 \end{bmatrix} * \begin{bmatrix} f(i - 1)\\ f(i - 2)\\ f(i - 3) \end{bmatrix}

我们令

Mat = \begin{bmatrix} 1 &1 &1 \\ 1 &0 &0 \\ 0 &1 &0 \end{bmatrix}

然后发现，利用 $Mat$ 我们也能实现数列递推（公式太难敲了，随便列两项吧）：

Mat * \begin{bmatrix} f(i - 1)\\ f(i - 2)\\ f(i - 3) \end{bmatrix} = \begin{bmatrix} f(i)\\ f(i - 1)\\ f(i - 2) \end{bmatrix}

Mat * \begin{bmatrix} f(i )\\ f(i - 1)\\ f(i - 2) \end{bmatrix} = \begin{bmatrix} f(i + 1)\\ f(i)\\ f(i - 1) \end{bmatrix}

再根据矩阵运算的结合律，最终有：

\begin{bmatrix} f(n)\\ f(n - 1)\\ f(n - 2) \end{bmatrix} = Mat^{n - 2} * \begin{bmatrix} f(2)\\ f(1)\\ f(0) \end{bmatrix}

从而将问题转化为求解 $Mat^{n - 2}$ ，这时候可以套用「矩阵快速幂」解决方案。

代码：

class Solution {
    int N = 3;
    int[][] mul(int[][] a, int[][] b) {
        int[][] c = new int[N][N];
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                c[i][j] = a[i][0] * b[0][j] + a[i][1] * b[1][j] + a[i][2] * b[2][j];
            }
        }
        return c;
    }
    public int tribonacci(int n) {
        if (n == 0) return 0;
        if (n == 1 || n == 2) return 1;
        int[][] ans = new int[][]{
            {1,0,0},
            {0,1,0},
            {0,0,1}
        };
        int[][] mat = new int[][]{
            {1,1,1},
            {1,0,0},
            {0,1,0}
        };
        int k = n - 2;
        while (k != 0) {
            if ((k & 1) != 0) ans = mul(ans, mat);
            mat = mul(mat, mat);
            k >>= 1;
        }
        return ans[0][0] + ans[0][1];
    }
}

时间复杂度： $O(\log{n})$
空间复杂度： $O(1)$

打表

当然，我们也可以将数据范围内的所有答案进行打表预处理，然后在询问时直接查表返回。

但对这种题目进行打表带来的收益没有平常打表题的大，因为打表内容不是作为算法必须的一个环节，而直接是作为该询问的答案，但测试样例是不会相同的，即不会有两个测试数据都是 $n = 37$ 。

这时候打表节省的计算量是不同测试数据之间的相同前缀计算量，例如 $n = 36$ 和 $n = 37$ ，其 $35$ 之前的计算量只会被计算一次。

因此直接为「解法二」的 cache 添加 static 修饰其实是更好的方式：代码更短，同时也能起到同样的节省运算量的效果。

代码：

class Solution {
    static int[] cache = new int[40];
    static {
        cache[0] = 0;
        cache[1] = 1;
        cache[2] = 1;
        for (int i = 3; i < cache.length; i++) {
            cache[i] = cache[i - 1] + cache[i - 2] + cache[i - 3];
        }
    }
    public int tribonacci(int n) {
        return cache[n];
    }
}

时间复杂度：将打表逻辑交给 $OJ$ ，复杂度为 $O(C)$ ， $C$ 固定为 $40$ 。将打表逻辑放到本地进行，复杂度为 $O(1)$
空间复杂度： $O(n)$

552. 学生出勤记录 II

可以用字符串表示一个学生的出勤记录，其中的每个字符用来标记当天的出勤情况（缺勤、迟到、到场）。

记录中只含下面三种字符：

'A'：Absent，缺勤
'L'：Late，迟到
'P'：Present，到场

如果学生能够同时满足下面两个条件，则可以获得出勤奖励：

按总出勤计，学生缺勤（'A'）严格少于两天。
学生不会存在连续 3 天或连续 3 天以上的迟到（'L'）记录。

给你一个整数 n ，表示出勤记录的长度（次数）。请你返回记录长度为 n 时，可能获得出勤奖励的记录情况数量。

答案可能很大，所以返回对 $10^9 + 7$ 取余的结果。

示例 1：

输入：n = 2

输出：8

解释：
有 8 种长度为 2 的记录将被视为可奖励：
"PP" , "AP", "PA", "LP", "PL", "AL", "LA", "LL" 
只有"AA"不会被视为可奖励，因为缺勤次数为 2 次（需要少于 2 次）。

提示：

1 <= n <= $10^5$

基本分析

根据题意，我们知道一个合法的方案中 A 的总出现次数最多为 $1$ 次，L 的连续出现次数最多为 $2$ 次。

因此在枚举/统计合法方案的个数时，当我们决策到某一位应该选什么时，我们关心的是当前方案中已经出现了多少个 A（以决策当前能否填入 A）以及连续出现的 L 的次数是多少（以决策当前能否填入 L）。

记忆化搜索

枚举所有方案的爆搜 DFS 代码不难写，大致的函数签名设计如下：

/**
 * @param u 当前还剩下多少位需要决策
 * @param acnt 当前方案中 A 的总出现次数
 * @param lcnt 当前方案中结尾 L 的连续出现次数
 * @param cur 当前方案
 * @param ans 结果集
 */
void dfs(int u, int acnt, int lcnt, String cur, List<String> ans);

实际上，我们不需要枚举所有的方案数，因此我们只需要保留函数签名中的前三个参数即可。

同时由于我们在计算某个 $(u, acnt, lcnt)$ 的方案数时，其依赖的状态可能会被重复使用，考虑加入记忆化，将结果缓存起来。

根据题意， $n$ 的取值范围为 $[0, n]$ ， $acnt$ 取值范围为 $[0,1]$ ， $lcnt$ 取值范围为 $[0, 2]$ 。

代码：

class Solution {
    int mod = (int)1e9+7;
    int[][][] cache;
    public int checkRecord(int n) {
        cache = new int[n + 1][2][3];
        for (int i = 0; i <= n; i++) {
            for (int j = 0; j < 2; j++) {
                for (int k = 0; k < 3; k++) {
                    cache[i][j][k] = -1;
                }
            }
        }
        return dfs(n, 0, 0);
    }
    int dfs(int u, int acnt, int lcnt) {
        if (acnt >= 2) return 0;
        if (lcnt >= 3) return 0;
        if (u == 0) return 1;
        if (cache[u][acnt][lcnt] != -1) return cache[u][acnt][lcnt];
        int ans = 0;
        ans = dfs(u - 1, acnt + 1, 0) % mod; // A
        ans = (ans + dfs(u - 1, acnt, lcnt + 1)) % mod; // L
        ans = (ans + dfs(u - 1, acnt, 0)) % mod; // P
        cache[u][acnt][lcnt] = ans;
        return ans;
    }
}

时间复杂度：有 $O(n * 2 * 3)$ 个状态需要被计算，复杂度为 $O(n)$
空间复杂度： $O(n)$

状态机 DP

通过记忆化搜索的分析我们发现，当我们在决策下一位是什么的时候，依赖于前面已经填入的 A 的个数以及当前结尾处的 L 的连续出现次数。

也就说是，状态 $f[u][acnt][lcnt]$ 必然被某些特定状态所更新，或者说由 $f[u][[acnt][lcnt]$ 出发，所能更新的状态是固定的。

因此这其实是一个状态机模型的 DP 问题。

根据「更新 $f[u][acnt][lcnt]$ 需要哪些状态值」还是「从 $f[u][acnt][lcnt]$ 出发，能够更新哪些状态」，我们能够写出两种方式（方向）的 DP 代码：

一类是从 $f[u][acnt][lcnt]$ 往回找所依赖的状态；一类是从 $f[u][acnt][lcnt]$ 出发往前去更新所能更新的状态值。

无论是何种方式（方向）的 DP 实现都只需搞清楚「当前位的选择对 $acnt$ 和 $lcnt$ 的影响」即可。

代码：

// 从 f[u][acnt][lcnt] 往回找所依赖的状态
class Solution {
    int mod = (int)1e9+7;
    public int checkRecord(int n) {
        int[][][] f = new int[n + 1][2][3];
        f[0][0][0] = 1;
        for (int i = 1; i <= n; i++) {
            for (int j = 0; j < 2; j++) {
                for (int k = 0; k < 3; k++) {
                    if (j == 1 && k == 0) { // A
                        f[i][j][k] = (f[i][j][k] + f[i - 1][j - 1][0]) % mod;
                        f[i][j][k] = (f[i][j][k] + f[i - 1][j - 1][1]) % mod;
                        f[i][j][k] = (f[i][j][k] + f[i - 1][j - 1][2]) % mod;
                    }
                    if (k != 0) { // L
                        f[i][j][k] = (f[i][j][k] + f[i - 1][j][k - 1]) % mod;
                    }
                    if (k == 0) { // P
                        f[i][j][k] = (f[i][j][k] + f[i - 1][j][0]) % mod;
                        f[i][j][k] = (f[i][j][k] + f[i - 1][j][1]) % mod;
                        f[i][j][k] = (f[i][j][k] + f[i - 1][j][2]) % mod;
                    }
                }
            }
        }
        int ans = 0;
        for (int j = 0; j < 2; j++) {
            for (int k = 0; k < 3; k++) {
                ans += f[n][j][k];
                ans %= mod;
            }
        }
        return ans;
    }
}

时间复杂度： $O(n)$
空间复杂度： $O(n)$

矩阵快速幂

之所以在动态规划解法中强调更新状态的方式（方向）是「往回」还是「往前」，是因为对于存在线性关系（同时又具有结合律）的递推式，我们能够通过「矩阵快速幂」来进行加速。

矩阵快速幂的基本分析之前在 (题解) 1137. 第 N 个泰波那契数详细讲过。

由于 $acnt$ 和 $lcnt$ 的取值范围都很小，其组合的状态只有 $2 * 3 = 6$ 种，我们使用 $idx = acnt * 3 + lcnt$ 来代指组合（通用的二维转一维方式）：

$idx = 0$ ： $acnt = 0、lcnt = 0$ ；
$idx = 1$ ： $acnt = 1、lcnt = 0$ ； ...
$idx = 5$ ： $acnt = 1、lcnt = 2$ ；

最终答案为 $ans = \sum_{idx = 0}^{5} f[n][idx]$ ，将答案依赖的状态整理成列向量：

g[n] = \begin{bmatrix} f[n][0]\\ f[n][1]\\ f[n][2]\\ f[n][3]\\ f[n][4]\\ f[n][5] \end{bmatrix}

根据状态机逻辑，可得：

g[n] = \begin{bmatrix} f[n][0]\\ f[n][1]\\ f[n][2]\\ f[n][3]\\ f[n][4]\\ f[n][5] \end{bmatrix} = \begin{bmatrix} f[n - 1][0] * 1 + f[n - 1][1] * 1 + f[n - 1][2] * 1 + f[n - 1][3] * 0 + f[n - 1][4] * 0 + f[n - 1][5] * 0\\ f[n - 1][0] * 1 + f[n - 1][1] * 0 + f[n - 1][2] * 0 + f[n - 1][3] * 0 + f[n - 1][4] * 0 + f[n - 1][5] * 0\\ f[n - 1][0] * 0 + f[n - 1][1] * 1 + f[n - 1][2] * 0 + f[n - 1][3] * 0 + f[n - 1][4] * 0 + f[n - 1][5] * 0\\ f[n - 1][0] * 1 + f[n - 1][1] * 1 + f[n - 1][2] * 1 + f[n - 1][3] * 1 + f[n - 1][4] * 1 + f[n - 1][5] * 1\\ f[n - 1][0] * 0 + f[n - 1][1] * 0 + f[n - 1][2] * 0 + f[n - 1][3] * 1 + f[n - 1][4] * 0 + f[n - 1][5] * 0\\ f[n - 1][0] * 0 + f[n - 1][1] * 0 + f[n - 1][2] * 0 + f[n - 1][3] * 0 + f[n - 1][4] * 1 + f[n - 1][5] * 0 \end{bmatrix}

我们令：

mat = \begin{bmatrix} 1 &1 &1 &0 &0 &0 \\ 1 &0 &0 &0 &0 &0 \\ 0 &1 &0 &0 &0 &0 \\ 1 &1 &1 &1 &1 &1 \\ 0 &0 &0 &1 &0 &0 \\ 0 &0 &0 &0 &1 &0 \end{bmatrix}

根据「矩阵乘法」即有：

g[n] = mat * g[n - 1]

起始时，我们只有 $g[0]$ ，根据递推式得：

g[n] = mat * mat * ... * mat * g[0]

再根据矩阵乘法具有「结合律」，最终可得：

g[n] = mat^n * g[0]

计算 $mat^n$ 可以套用「快速幂」进行求解。

代码：

class Solution {
    int N = 6;
    int mod = (int)1e9+7;
    long[][] mul(long[][] a, long[][] b) {
        int r = a.length, c = b[0].length, z = b.length;
        long[][] ans = new long[r][c];
        for (int i = 0; i < r; i++) {
            for (int j = 0; j < c; j++) {
                for (int k = 0; k < z; k++) {
                    ans[i][j] += a[i][k] * b[k][j];
                    ans[i][j] %= mod;
                }
            }
        }
        return ans;
    }
    public int checkRecord(int n) {
        long[][] ans = new long[][]{
            {1}, {0}, {0}, {0}, {0}, {0}
        };
        long[][] mat = new long[][]{
            {1, 1, 1, 0, 0, 0},
            {1, 0, 0, 0, 0, 0},
            {0, 1, 0, 0, 0, 0},
            {1, 1, 1, 1, 1, 1},
            {0, 0, 0, 1, 0, 0},
            {0, 0, 0, 0, 1, 0}
        };
        while (n != 0) {
            if ((n & 1) != 0) ans = mul(mat, ans);
            mat = mul(mat, mat);
            n >>= 1;
        }
        int res = 0;
        for (int i = 0; i < N; i++) {
            res += ans[i][0];
            res %= mod;
        }
        return res;
    } 
}

时间复杂度： $O(\log{n})$
空间复杂度： $O(1)$

剑指 Offer 10- I. 斐波那契数列

写一个函数，输入 n ，求斐波那契（Fibonacci）数列的第 n 项（即 F(N)）。斐波那契数列的定义如下：

F(0) = 0, F(1) = 1
F(N) = F(N - 1) + F(N - 2), 其中 N > 1.

斐波那契数列由 0 和 1 开始，之后的斐波那契数就是由之前的两数相加而得出。

答案需要取模 1e9+7（1000000007），如计算初始结果为：1000000008，请返回 1。

示例 1：

输入：n = 2

输出：1

提示：

0 <= n <= 100

递推实现动态规划

既然转移方程都给出了，直接根据转移方程从头到尾递递推一遍即可。

代码：

class Solution {
    int mod = (int)1e9+7;
    public int fib(int n) {
        if (n <= 1) return n;
        int a = 0, b = 1;
        for (int i = 2; i <= n; i++) {
            int c = a + b;
            c %= mod;
            a = b;
            b = c;
        }
        return b;
    }
}

时间复杂度： $O(n)$
空间复杂度： $O(1)$

递归实现动态规划

能以「递推」形式实现动态规划，自然也能以「递归」的形式实现。

为防止重复计算，我们需要加入「记忆化搜索」功能，同时利用某个值 $x$ 在不同的样例之间可能会作为“中间结果”被重复计算，并且计算结果 $fib(x)$ 固定，我们可以使用 static 修饰缓存器，以实现计算过的结果在所有测试样例中共享。

代码：

class Solution {
    static int mod = (int)1e9+7;
    static int N = 110;
    static int[] cache = new int[N];
    public int fib(int n) {
        if (n <= 1) return n;
        if (cache[n] != 0) return cache[n];
        cache[n] = fib(n - 1) + fib(n - 2);
        cache[n] %= mod;
        return cache[n];
    }
}

时间复杂度： $O(n)$
空间复杂度： $O(1)$

打表

经过「解法二」，我们进一步发现，可以利用数据范围只有 $100$ 进行打表预处理，然后直接返回。

代码：

class Solution {
    static int mod = (int)1e9+7;
    static int N = 110;
    static int[] cache = new int[N];
    static {
        cache[1] = 1;
        for (int i = 2; i < N; i++) {
            cache[i] = cache[i - 1] + cache[i - 2];
            cache[i] %= mod;
        }
    }
    public int fib(int n) {
        return cache[n];
    }
}

时间复杂度：将打表逻辑放到本地执行，复杂度为 $O(1)$ ；否则为 $O(C)$ ， $C$ 为常量，固定为 $110$
空间复杂度： $O(C)$

矩阵快速幂

对于数列递推问题，可以使用矩阵快速幂进行加速，最完整的介绍在这里讲过。

对于本题，某个 $f(n)$ 依赖于 $f(n - 1)$ 和 $f(n - 2)$ ，将其依赖的状态存成列向量：

\begin{bmatrix} f(n - 1)\\ f(n - 2) \end{bmatrix}

目标值 $f(n)$ 所在矩阵为：

\begin{bmatrix} f(n)\\ f(n - 1) \end{bmatrix}

根据矩阵乘法，不难发现：

\begin{bmatrix} f(n)\\ f(n - 1) \end{bmatrix} = \begin{bmatrix} 1& 1\\ 1& 0 \end{bmatrix} * \begin{bmatrix} f(n - 1)\\ f(n - 2) \end{bmatrix}

我们令：

mat = \begin{bmatrix} 1& 1\\ 1& 0 \end{bmatrix}

起始时，我们只有 $\begin{bmatrix} f(1)\\ f(0) \end{bmatrix}$ ，根据递推式得：

\begin{bmatrix} f(n)\\ f(n - 1) \end{bmatrix} = mat * mat * ... * mat * \begin{bmatrix} f(1)\\ f(0) \end{bmatrix}

再根据矩阵乘法具有「结合律」，最终可得：

\begin{bmatrix} f(n)\\ f(n - 1) \end{bmatrix} = mat^{n - 1} * \begin{bmatrix} f(1)\\ f(0) \end{bmatrix}

计算 $mat^{n - 1}$ 可以套用「快速幂」进行求解。

代码：

class Solution {
    int mod = (int)1e9+7;
    long[][] mul(long[][] a, long[][] b) {
        int r = a.length, c = b[0].length, z = b.length;
        long[][] ans = new long[r][c];
        for (int i = 0; i < r; i++) {
            for (int j = 0; j < c; j++) {
                for (int k = 0; k < z; k++) {
                    ans[i][j] += a[i][k] * b[k][j];
                    ans[i][j] %= mod;
                }
            }
        }
        return ans;
    }
    public int fib(int n) {
        if (n <= 1) return n;
        long[][] mat = new long[][]{
            {1, 1},
            {1, 0}
        };
        long[][] ans = new long[][]{
            {1},
            {0}
        };
        int x = n - 1;
        while (x != 0) {
            if ((x & 1) != 0) ans = mul(mat, ans);
            mat = mul(mat, mat);
            x >>= 1;
        }
        return (int)(ans[0][0] % mod);
    }
}

时间复杂度： $O(\log{n})$
空间复杂度： $O(1)$

总结

综上，对于此类的「数列递推」问题，我们可以使用「矩阵快速幂」来进行加速（比如要递归一个长度为 $1e9$ 的数列，线性复杂度会被卡）。

对于转移存在线性关系（同时又具有结合律）的递推式，我们能够通过「矩阵快速幂」来进行加速。

有效将复杂度从 $O(n)$ 优化至 $O(\log{n})$ 。