三步问题【python，算法，leetcode】

爬楼梯问题：小明一次可以爬 1 个楼梯、2 个楼梯、3 楼梯，问要爬上 n 阶楼梯，小明可以有多少中爬法？

解决方案：使用矩阵的幂次方可以快速计算出爬 n 阶楼梯总共有多少种爬法。

假设f(n)表示爬n阶楼梯总共的方案数，其中f(2)=2、f(1）=1、f(0)=1,那么可以得到以下通用公式：
$${\left[\begin{array}{ccc}1& 1& 1\\ 1& 0& 0\\ 0& 1& 0\end{array}\right]}^{n-2}\left[\begin{array}{c}f(2)\\ f(1)\\ f(0)\end{array}\right]=\left[\begin{array}{c}f(n)\\ f(n-1)\\ f(n-2)\end{array}\right]$$

根据这个公式我们用 Python 进行代码实现。

import numpy as np


# 求矩阵的幂次方
def get_matrix_power(mat, power, mod):
	if power == 0:
		return np.mat(np.eye(mat.shape[0]))
	if power == 1:
		return mat
	if power % 2 == 0:
		return get_matrix_power(mat, power // 2, mod) ** 2 % mod
	else:
		return mat * get_matrix_power(mat, power - 1, mod) % mod


# 计算爬楼梯总统有多少中方法
def ways_to_steps(n):
	# 对结果需要取模，防止计算溢出
	mod = 1000000007
	# 初始值
	ans = [1, 1, 2]
	if n <= 2:
		return ans[n]
	# 定义基本矩阵
	mat = np.mat([[1, 1, 1],
	              [1, 0, 0],
	              [0, 1, 0]])
	# 对矩阵进行幂次方
	power_base_mat = get_matrix_power(mat, n - 2, mod)
	# 定义初始函数矩阵
	base_f_mat = np.mat([[2], [1], [1]])
	# 两个矩阵相乘，得到结果
	res = (power_base_mat * base_f_mat % mod)[0][0]
	return int(res)

当然，计算矩阵的n次幂还有很多方法，你可以根据实际需要进行修改。

注意：题目源自 leetcode，链接为 https://leetcode.cn/problems/three-steps-problem-lcci/description/