面试必考的「矩阵快速幂」考点汇总

大家好，欢迎来到IT知识分享网。

设想这样一个场景，面试官给了你一道算法题，你很快确定这是一道递推问题，并给出了 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 阶。实现一种方法，计算小孩有多少种上楼梯的方式。结果可能很大，你需要对结果模 。

示例 1

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

示例 2

输入：n = 5 输出：13

解题思路

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

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

接下来确定矩阵

首先假设

中仅包含

与

由于我们需要用

推出

根据之前的假设，

中仅包含

和

而推出

却需要

和

因此假设不成立，即

中应该包含

即

​

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

由于求解

只需求出

因此

中包含

根据题意可得：

由此使用快速幂求出

再与

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

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

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 %  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]) %  return c

总结

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

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

确定转移矩阵 B

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

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

本文作者：Gene_Liu

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