设想这样一个场景，面试官给了你一道算法题，你很快确定这是一道递推问题，并给出了 O(n) 的解题方法，然而面试官却继续问：“还能继续优化吗？”

这样类似的场景并不少见，因为算法不仅追求「正确」，还追求「效率」，而这也正是优化方法的意义。

本文即将介绍的「矩阵快速幂」便是一种常见的优化递推的方法，能将 O(n) 的递推过程加速至 O(log(n))，使效率显著提升。

一、矩阵运算

首先我们回顾一下「矩阵」，一个 n * m 的矩阵可以看作是一个 n * m 的二维数组，而矩阵的加减法就是将两个矩阵对应位置上的数相加减，即 C = A + B 意味着矩阵 C 中任意一点

其中 A, B, C 均为 n * m 的矩阵，具体例子如下：

矩阵乘法的运算稍微复杂一些，假设 A 是 n * m 的矩阵，B 是 m * p 的矩阵，则 C = A * B 是 n * p 的矩阵，且对于 C 中任意一点

来说，满足：

即 C 中第 i 行、第 j 列的元素等于 A 中第 i 行与 B 中第 j 列所有元素对应相乘再相加，这也就意味着矩阵相乘时，第一个矩阵的列数必须等于第二个矩阵的行数，具体例子如下：

另外，矩阵乘法满足结合律，即 (A * B) * C = A * (B * C)，举例如下：

满足分配律，即 (A + B) * C = A * C + B * C，具体例子如下：

但不满足交换律，即 A * B 不一定等于 B * A，具体例子如下：

代码实现

我们将一个 n * m 的矩阵看作是一个 n * m 的二维数组，因此在 C++ 中我们用 vector<vector<int>> 来表示，矩阵相乘函数如下：

vector<vector<int>> matrix_multiply(vector<vector<int>>& a, vector<vector<int>>& b) {
    int n = a.size(), m = a[0].size(), p = b[0].size();
    vector<vector<int>> c(n, vector<int>(p, 0));
    for (int i = 0; i < n; i++)
        for (int j = 0; j < p; j++)
            for (int k = 0; k < m; k++)
                c[i][j] += a[i][k] * b[k][j];
    return c;
}

在 Python3 中我们用 List[List[int]] 来表示，矩阵相乘函数如下：

def matrix_multiply(a: List[List[int]], b: List[List[int]]) -> List[List[int]]:
    n, m, p = len(a), len(a[0]), len(b[0])
    c = [[0 for j in range(p)] for i in range(n)]
    for i in range(n):
        for j in range(p):
            for k in range(m):
                c[i][j] += a[i][k] * b[k][j]
    return c

二、快速幂

讲解「矩阵快速幂」前，我们需要回顾一下「快速幂」，完全没有了解过的同学建议翻阅「刷算法题必备的数学考点汇总 」中相关内容。

简单来说，「快速幂」就是通过将指数转化为二进制，并以此来加快幂运算，例如 2^31 中 31 满足：

其中 D 表示十进制，B 表示二进制。进一步地，2^31 的计算过程可以转变为：

因此我们只需从 1 开始计算 5 次求出

再将它们依次相乘，即可得到 2^31。

当计算 a^n（a 为任意实数）时，快速幂方法可以将原来的 O(n) 时间复杂度降低为 O(log(n))，从而大大加快指数运算。

三、矩阵快速幂

接下来开始「矩阵快速幂」的介绍，首先以经典「斐波那契数列」为例进行讲解。

509. 斐波那契数

题目描述

斐波那契数，通常用 F(n) 表示，形成的序列称为斐波那契数列 。该数列由 0 和 1 开始，后面的每一项数字都是前面两项数字的和。也就是：

F(0) = 0，F(1) = 1
F(n) = F(n - 1) + F(n - 2)，其中 n > 1

给你 n ，请计算 F(n) 。

示例 1

输入：2
输出：1
解释：F(2) = F(1) + F(0) = 1 + 0 = 1

示例 2

输入：3
输出：2
解释：F(3) = F(2) + F(1) = 1 + 1 = 2

示例 3

输入：4
输出：3
解释：F(4) = F(3) + F(2) = 2 + 1 = 3

解题思路

此题是经典的递推问题，即

我们可以很轻松地在 O(n) 时间复杂度内求出

但假如 n 很大呢？

当 n = 1e9 时，显然单纯的递推无法通过此题，因此我们用矩阵快速幂来加速递推过程。

首先我们需要构造矩阵，令

则从

到

的转移矩阵 B 为：

进一步地，我们可以得到：

由于矩阵乘法满足结合律，因此矩阵 B 的幂次运算也可以通过「快速幂」进行加速，例如：

因此本题我们可以通过快速计算

再将其与矩阵

相乘，即可得到

将原来 O(n) 的递推加速至 O(log(n)) 的这一过程即为「矩阵快速幂」。

另外，由于

中已包含

因此只需计算

具体细节见代码。

C++ 代码

class Solution {
public:
    int fib(int n) {
        if (n < 2) return n;
        vector<vector<int>> q{{1, 1}, {1, 0}};
        vector<vector<int>> res = matrix_pow(q, n - 1);
        return res[0][0];
    }

    vector<vector<int>> matrix_pow(vector<vector<int>>& a, int n) {
        vector<vector<int>> ret{{1, 0}, {0, 1}};
        while (n > 0) {
            if (n & 1) ret = matrix_multiply(ret, a);
            n >>= 1;
            a = matrix_multiply(a, a);
        }
        return ret;
    }

    vector<vector<int>> matrix_multiply(vector<vector<int>>& a, vector<vector<int>>& b) {
        int n = a.size(), m = a[0].size(), p = b[0].size();
        vector<vector<int>> c(n, vector<int>(p, 0));
        for (int i = 0; i < n; i++)
            for (int j = 0; j < p; j++)
                for (int k = 0; k < m; k++)
                    c[i][j] += a[i][k] * b[k][j];
        return c;
    }
};

Python3 代码

class Solution:
    def fib(self, n: int) -> int:
        if n < 2:
            return n
        q = [[1, 1], [1, 0]]
        res = self.matrix_pow(q, n - 1)
        return res[0][0]
    
    def matrix_pow(self, a: List[List[int]], n: int) -> List[List[int]]:
        ret = [[1, 0], [0, 1]]
        while n > 0:
            if n & 1:
                ret = self.matrix_multiply(ret, a)
            n >>= 1
            a = self.matrix_multiply(a, a)
        return ret

    def matrix_multiply(self, a: List[List[int]], b: List[List[int]]) -> List[List[int]]:
        n, m, p = len(a), len(a[0]), len(b[0])
        c = [[0 for j in range(p)] for i in range(n)]
        for i in range(n):
            for j in range(p):
                for k in range(m):
                    c[i][j] += a[i][k] * b[k][j]
        return c

运用「矩阵快速幂」加速递推过程时，首先需要确定递推公式，然后再根据递推公式，确定矩阵

中究竟是

还是

或者包含更多。接下来再根据

和

确定转移矩阵 B，最后应用快速幂计算 B 的幂次即可完成求解。

接下来我们再通过一道例题加强对该解题思路的掌握。

面试题 08.01. 三步问题

题目描述

三步问题。有个小孩正在上楼梯，楼梯有 n 阶台阶，小孩一次可以上 1 阶、2 阶或 3 阶。实现一种方法，计算小孩有多少种上楼梯的方式。结果可能很大，你需要对结果模 1000000007。

示例 1

输入：n = 3 
输出：4
说明: 有四种走法

示例 2

输入：n = 5
输出：13

解题思路

首先确定本题的递推公式，令

表示上 n 阶台阶的方法数，则不难发现

接下来确定矩阵

首先假设

中仅包含

与

由于我们需要用

推出

根据之前的假设，

中仅包含

和

而推出

却需要

和

因此假设不成立，即

中应该包含

即

由此我们可以确定转移矩阵 B，即

由于求解

只需求出

因此

中包含

根据题意可得：

由此使用快速幂求出

再与

相乘即可在 O(log(n)) 的时间复杂度内求出

另外本题需要对结果模 1000000007，即在矩阵乘法与最终答案计算时进行取模，具体细节见下述代码。

C++ 代码

class Solution {
public:
    typedef long long ll;
    static const int mod = 1e9 + 7;
    int waysToStep(int n) {
        if (n < 3) return n;
        vector<vector<int>> q{{1, 1, 1}, {1, 0, 0}, {0, 1, 0}};
        vector<vector<int>> res = matrix_pow(q, n - 2);
        ll ans = res[0][0] * 2ll + res[0][1] + res[0][2];
        return ans % mod;
    }

    vector<vector<int>> matrix_pow(vector<vector<int>>& a, int n) {
        vector<vector<int>> ret = {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}};
        while (n > 0) {
            if (n & 1) ret = matrix_multiply(ret, a);
            n >>= 1;
            a = matrix_multiply(a, a);
        }
        return ret;
    }

    vector<vector<int>> matrix_multiply(vector<vector<int>>& a, vector<vector<int>>& b) {
        int n = a.size(), m = a[0].size(), p = b[0].size();
        vector<vector<int>> c(n, vector<int>(p, 0));
        for (int i = 0; i < n; i++)
            for (int j = 0; j < p; j++)
                for (int k = 0; k < m; k++)
                    c[i][j] = (c[i][j] + (ll)a[i][k] * (ll)b[k][j] % mod) % mod;
        return c;
    }
};

Python3 代码

class Solution:
    def waysToStep(self, n: int) -> int:
        if n < 3:
            return n
        q = [[1, 1, 1], [1, 0, 0], [0, 1, 0]]
        res = self.matrix_pow(q, n - 2)
        ans = res[0][0] * 2 + res[0][1] + res[0][2]
        return ans % 1000000007
    
    def matrix_pow(self, a: List[List[int]], n: int) -> List[List[int]]:
        ret = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
        while n > 0:
            if n & 1:
                ret = self.matrix_multiply(ret, a)
            n >>= 1
            a = self.matrix_multiply(a, a)
        return ret

    def matrix_multiply(self, a: List[List[int]], b: List[List[int]]) -> List[List[int]]:
        n, m, p = len(a), len(a[0]), len(b[0])
        c = [[0 for j in range(p)] for i in range(n)]
        for i in range(n):
            for j in range(p):
                for k in range(m):
                    c[i][j] = (c[i][j] + a[i][k] * b[k][j]) % 1000000007
        return c

总结

「矩阵快速幂」是一种常见的递推优化方法，可以将 O(n) 的递推优化至 O(log(n))，具体算法流程如下：

1. 确定递推式
2. 确定矩阵 An 的维数
3. 根据 An 和

确定转移矩阵 B

4.快速幂计算转移矩阵 B 的幂次

5.将 B 的幂次运算结果与A0 相乘得到最终答案 Fn

本文作者：Gene_Liu

声明：本文归 “力扣” 版权所有，如需转载请联系。