[LeetCode] Climbing Stairs

Note: If you feel unwilling to read the long codes, just take the idea with you. The codes are unnecessarily long due to the inconvenient handle of matrices.

Well, a classic and interesting problem. The recursion is simply f(n) = f(n - 1) + f(n - 2), which means that we can either climb to n - 1 and then climb 1 step or climb to n - 2 and then climb 2 steps. So this problem is actually asking for the n-th Fibonacci number. However, if you code it in such a recursive way, it will meet TLE due to the large number of overlapping sub-problems.

There are mainly two ways to solve this problem. The first one uses the above formula in a bottom-up manner and takes O(n) time. This link shares the O(n) solution in all the supported languages of the LeetCode OJ. You may take a look at it and appreciate the appetite of each language :-)

Now I will focus on another solution, which takes O(logn) time. The idea is to use the matrix power. In fact, [f(n), f(n - 1); f(n - 1), f(n - 2)] = [1, 1; 1, 0] ^ n for n >= 2. And similar to the problem Pow(x, n), the power of a matrix can be computed in O(logn) time.

The C++ and Python codes are shown below. Note that the computation of the power of the matrix[1, 1; 1, 0] is hard-coded. Since it is a bit trickier to handle matrix multiplications, the codes become much longer.

C++

 1 class Solution {  
 2 public:
 3     int climbStairs(int n) {
 4         if (n < 2) return n;
 5         vector<int> fibs = {1, 1, 1, 0};
 6         vector<int> ans = fibPower(fibs, n);
 7         return ans[0];
 8     }
 9 private:
10     vector<int> matrixProd(vector<int>& l, vector<int>& r) {
11         vector<int> ans(4, 0);
12         ans[0] = l[0] * r[0] + l[1] * r[2];
13         ans[1] = l[0] * r[1] + l[1] * r[3];
14         ans[2] = l[2] * r[0] + l[3] * r[2];
15         ans[3] = l[2] * r[1] + l[3] * r[3]; 
16         return ans;
17     }
18     vector<int> fibPower(vector<int>& fibs, int n){
19         if (n == 1) return fibs;
20         vector<int> half1 = fibPower(fibs, n / 2);
21         vector<int> half2 = fibPower(fibs, n / 2);
22         vector<int> ans = matrixProd(half1, half2);
23         if (n % 2 == 0) return ans;
24         ans[1] = (ans[0] += ans[1]) - ans[1];
25         ans[3] = (ans[2] += ans[3]) - ans[3];
26         return ans;
27     }
28 };

Python

class Solution:
    # @param {integer} n
    # @return {integer}
    def climbStairs(self, n):
        if n  < 2:
            return n
        fibs = [1, 1, 1, 0]
        ans = self.fibsPower(fibs, n)
        return ans[0] 

    def matrixProd(self, l, r):
        ans = [0] * 4
        ans[0] = l[0] * r[0] + l[1] * r[2]
        ans[1] = l[0] * r[1] + l[1] * r[3]
        ans[2] = l[2] * r[0] + l[3] * r[2]
        ans[3] = l[2] * r[1] + l[3] * r[3]
        return ans

    def fibsPower(self, fibs, n):
        if n == 1:
            return fibs
        half1 = self.fibsPower(fibs, n / 2)
        half2 = self.fibsPower(fibs, n / 2)
        ans = self.matrixProd(half1, half2) 
        if n % 2 == 0:
            return ans
        ans[0], ans[1], ans[2], ans[3] = ans[0] + ans[1], ans[0], ans[2] + ans[3], ans[2]
        return ans