2019校招真题在线编程 爬楼梯2（矩阵快速幂）

题目描述

在你面前有一个n阶的楼梯(n>=100且n<500)，你一步只能上1阶或3阶。
请问计算出你可以采用多少种不同的方式爬完这个楼梯（到最后一层为爬完）。

输入描述:

一个正整数，表示这个楼梯一共有多少阶

输出描述:

一个正整数，表示有多少种不同的方式爬完这个楼梯

示例1

输入

复制

100

输出

复制

24382819596721629

题解

    =  *

def matmul(A, B):
    C = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
    for i in range(3):
        for j in range(3):
            for k in range(3):
                C[i][j] += (A[i][k] * B[k][j])
    return C
 
def mat_power(A, count):
    B = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    while (count > 0) :
        if count % 2 == 1:
            B = matmul(B, A)
        A = matmul(A, A)
        count //= 2
    return B
 
def work():
    n = int(input())
    f = [0 for i in range(n + 1)]
    f[0] = f[1] = f[2] = 1
    f[3] = 2
    if n <= 3:
        print(f[n])
        return
    A = [[f[2], f[0], f[1]], [f[3], f[1], f[2]], [0, 0, 0]]
    B = [[1, 0, 1], [1, 0, 0], [0, 1, 0]]
    count = n - 2 - n % 2
    B = mat_power(B, count)
    B = matmul(A, B)
    if n % 2 == 1:
        print(B[1][0])
    else:
        print(B[0][0])
work()