难度: Easy
三步问题。有个小孩正在上楼梯，楼梯有n阶台阶，小孩一次可以上1阶、2阶或3阶。实现一种方法，计算小孩有多少种上楼梯的方式。结果可能很大，你需要对结果模1000000007。
示例1:
输入：n = 3 输出：4 说明: 有四种走法
示例2:
输入：n = 5 输出：13
提示:
- n范围在[1, 1000000]之间
Submission
运行时间: 38 ms
内存: 16.2 MB
def matrix_multiply(a, b, MOD=10 ** 9 + 7):
m, n, p = len(a), len(a[0]), len(b[0])
ans = [[0] * p for _ in range(m)]
for i in range(m):
for j in range(n):
for k in range(p): # 乘法也可以优化成分治法
ans[i][k] += (a[i][j] * b[j][k]) % MOD
ans[i][k] %= MOD
return ans
def matrix_pow_mod(a, b, MOD=10 ** 9 + 7):
n = len(a)
ans = [[0] * n for _ in range(n)]
for i in range(n):
ans[i][i] = 1
while b:
if b & 1:
ans = matrix_multiply(ans, a, MOD)
a = matrix_multiply(a, a, MOD)
b >>= 1
return ans
class Solution:
def waysToStep(self, n: int) -> int:
if n == 0: return 0
m = [[1, 1, 1], [1, 0, 0], [0, 1, 0]]
return matrix_pow_mod(m, n)[0][0]
Explain
本题的解决方案采用矩阵快速幂来解决上楼梯问题。首先将问题转化为状态转移方程，其中每个状态表示从前一级、前两级或前三级台阶达到当前台阶的方式。接着，使用一个转移矩阵将这种状态转移表达出来，然后通过矩阵的快速幂求解来获得第n阶的计算结果。矩阵的快速幂使得我们可以在对数时间内得到结果，避免了线性时间上的计算。
时间复杂度: O(log n)
空间复杂度: O(1)
def matrix_multiply(a, b, MOD=10 ** 9 + 7):
m, n, p = len(a), len(a[0]), len(b[0])
ans = [[0] * p for _ in range(m)]
for i in range(m):
for j in range(n):
for k in range(p): # 乘法也可以优化成分治法
ans[i][k] += (a[i][j] * b[j][k]) % MOD
ans[i][k] %= MOD
return ans
def matrix_pow_mod(a, b, MOD=10 ** 9 + 7):
n = len(a)
ans = [[0] * n for _ in range(n)]
for i in range(n):
ans[i][i] = 1 # 初始化单位矩阵
while b: # 利用快速幂算法减少乘法次数
if b & 1: # 如果最低位为1，累乘当前矩阵
ans = matrix_multiply(ans, a, MOD)
a = matrix_multiply(a, a, MOD) # 矩阵自乘，加速幂的计算
b >>= 1 # 位移，减少一半的幂次
return ans
class Solution:
def waysToStep(self, n: int) -> int:
if n == 0: return 0
m = [[1, 1, 1], [1, 0, 0], [0, 1, 0]] # 转移矩阵定义
return matrix_pow_mod(m, n)[0][0] # 计算结果，返回第一行第一列的值，即为方法数
Explore
在矩阵快速幂的实现过程中，为了避免整数溢出，每次矩阵乘法计算时，都会对结果进行模运算 `MOD = 10**9 + 7`。这个模运算是在乘法累加的每一步之后立即进行的，确保每个中间结果都不会超过模数，从而防止了在后续计算中的整数溢出。
状态转移方程源于解决上楼梯问题，其中可以从前一级、前两级或前三级台阶上到当前台阶。定义状态 dp[i] 为到达第 i 级台阶的方法数量。状态转移方程为 dp[i] = dp[i-1] + dp[i-2] + dp[i-3]。在矩阵表达中，这转化为使用一个转移矩阵 m = [[1, 1, 1], [1, 0, 0], [0, 1, 0]]，其中每个元素 m[i][j] 表示 dp[j] 对 dp[i] 的贡献。
最终结果只需返回矩阵乘法结果中的第一行第一列的值，因为这个值代表从基状态 (即 dp[1], dp[2], dp[3]) 到达第 n 级台阶的方法总数。在矩阵 m 的乘幂表示中，第一行的元素累积了所有可能的到达第 n 级台阶的路径，因此 matrix_pow_mod(m, n)[0][0] 正是我们需要的达到第 n 级台阶的方式数。