题解 | 上高楼_牛客博客

解题思路

这是一个斐波那契数列的变种问题，使用矩阵快速幂来解决。关键点：

递推关系：
- $f(n) = f(n-1) + f(n-2)$
- 可以使用矩阵快速幂优化时间复杂度
矩阵转换：
- 使用 $2\times 2$ 矩阵进行状态转移
- 注意矩阵下标从 $1$ 开始
取模运算：
- 所有运算都需要取模，防止溢出
- mod = 1000000007

代码

cpp
java
python

class GoUpstairs {
public:
    #define i64 long long
    const int mod = 1e9+7;
    
    vector<vector<i64>> ksm_mar(vector<vector<i64>>A, vector<vector<i64>>B) {
        vector<vector<i64>>C(3, vector<i64>(3, 0));
        for(int i = 1; i <= 2; i++) {
            for(int j = 1; j <= 2; j++) {
                for(int k = 1; k <= 2; k++) {
                    C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % mod) % mod;
                }
            }
        }
        return C;
    }
    
    vector<vector<i64>> cul(vector<vector<i64>>A, vector<vector<i64>>B) {
        vector<vector<i64>>C(3, vector<i64>(3, 0));
        for(int i = 1; i <= 2; i++) {
            for(int j = 1; j <= 1; j++) {
                for(int k = 1; k <= 2; k++) {
                    C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % mod) % mod;
                }
            }
        }
        return C;
    }
    
    int countWays(int n) {
        if(n == 0) return 1;
        
        vector<vector<i64>>A(3, vector<i64>(3, 0));
        A[1][2] = 1;
        A[2][1] = 1;
        A[2][2] = 1;
        
        vector<vector<i64>>ans(3, vector<i64>(3, 0));
        ans[1][1] = 1;
        ans[2][1] = 2;
        
        n--;
        while(n) {
            if(n & 1) ans = cul(A, ans);
            A = ksm_mar(A, A);
            n >>= 1;
        }
        return ans[1][1];
    }
};

import java.util.*;

public class GoUpstairs {
    private static final int mod = 1000000007;
    
    private long[][] ksm_mar(long[][] A, long[][] B) {
        long[][] C = new long[3][3];
        for(int i = 1; i <= 2; i++) {
            for(int j = 1; j <= 2; j++) {
                for(int k = 1; k <= 2; k++) {
                    C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % mod) % mod;
                }
            }
        }
        return C;
    }
    
    private long[][] cul(long[][] A, long[][] B) {
        long[][] C = new long[3][3];
        for(int i = 1; i <= 2; i++) {
            for(int j = 1; j <= 1; j++) {
                for(int k = 1; k <= 2; k++) {
                    C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % mod) % mod;
                }
            }
        }
        return C;
    }
    
    public int countWays(int n) {
        if(n == 0) return 1;
        
        long[][] A = new long[3][3];
        A[1][2] = 1;
        A[2][1] = 1;
        A[2][2] = 1;
        
        long[][] ans = new long[3][3];
        ans[1][1] = 1;
        ans[2][1] = 2;
        
        n--;
        while(n > 0) {
            if((n & 1) == 1) ans = cul(A, ans);
            A = ksm_mar(A, A);
            n >>= 1;
        }
        return (int)ans[1][1];
    }
}

class GoUpstairs:
    def ksm_mar(self, A, B):
        mod = 1000000007
        C = [[0 for i in xrange(3)] for i in xrange(3)]
        for i in xrange(1, 3):
            for j in xrange(1, 3):
                for k in xrange(1, 3):
                    C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % mod) % mod
        return C
    
    def cul(self, A, B):
        mod = 1000000007
        C = [[0 for i in xrange(3)] for i in xrange(3)]
        for i in xrange(1, 3):
            for j in xrange(1, 2):
                for k in xrange(1, 3):
                    C[i][j] = (C[i][j] + (A[i][k] * B[k][j]) % mod) % mod
        return C
    
    def countWays(self, n):
        if n == 0:
            return 1
            
        mod = 1000000007
        A = [[0 for i in xrange(3)] for i in xrange(3)]
        A[1][2] = 1
        A[2][1] = 1
        A[2][2] = 1
        
        ans = [[0 for i in xrange(3)] for i in xrange(3)]
        ans[1][1] = 1
        ans[2][1] = 2
        
        n -= 1
        while n:
            if n & 1:
                ans = self.cul(A, ans)
            A = self.ksm_mar(A, A)
            n >>= 1
        
        return ans[1][1]

算法及复杂度

算法：矩阵快速幂
时间复杂度： $\mathcal{O(\log n)}$ ，使用快速幂优化
空间复杂度： $\mathcal{O(1)}$ ，只使用了常数额外空间