设想这样一个场景，面试官给了你一道算法题，你很快确定这是一道递推问题，并给出了 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))，具体算法流程如下：
- 确定递推式
- 确定矩阵 An 的维数
- 根据 An 和
确定转移矩阵 B
4.快速幂计算转移矩阵 B 的幂次
5.将 B 的幂次运算结果与A0 相乘得到最终答案 Fn
本文作者：Gene_Liu
声明：本文归 “力扣” 版权所有，如需转载请联系。