Using only 1, 2, 3 steps reach to nth step

Question

I was doing an online interview some days back and came across this task that we can reach the nth step by using only steps 1 or 2 or 3.

Updated Question: A child is running up a staircase with n steps and can hop either 1 step, 2 steps, or 3 steps at a time. Implement a method to count how many possible ways the child can run up the stairs.

This is my solution using dynamic programming. But during the interview, I saw that there were 6 hidden test cases that were failing.

public static long countWays(int n)
    {
        if(n==0 || n==1) return 1;

        long[] res = new long[n + 1];
        res[0] = 1;
        res[1] = 1;
        res[2] = 2;
 
        for (int i = 3; i <= n; i++)
            res[i] = res[i - 1] + res[i - 2]
                     + res[i - 3];
 
        return res[n];
    }
 

During the interview, only 5 test cases were passing out of 11. The remaining 6 are hidden test cases.

The range for n is 0 to 10 power 8.

I was not able to understand where I was doing a mistake, already this code is in O(n) time complexity. Is there a better way to do this or am I missing something?

Answer 1

You can write down the recurrence relation in this problem

S_k = S_k-1 + S_k-2 + S_k-3

in the following matrix form:

| S[k]   |   | 1  1  1 |   | S[k-1] |   
| S[k-1] | = | 1  0  0 | * | S[k-2] |
| S[k-2] |   | 0  1  0 |   | S[k-3] |   

Which means that the k-th power of the matrix can quickly give you S_k:

            k
| 1  1  1 |     | S[2] |   | S[k+2] | 
| 1  0  0 |   * | S[1] | = | S[k+1] |
| 0  1  0 |     | S[0] |   | S[k]   |

So to solve this problem you basically need to find the n-th power of this matrix.

You can raise something to the n-th power in O(log n) time using the exponentiation by squaring algorithm.

Another useful observation is that you need only 3 numbers S_k, S_k-1 and S_k-2 to the represent the k-th power of the matrix:

            k
| 1  1  1 |     | S[k] + S[k-1] + S[k-2]    S[k] + S[k-1]      S[k]   |
| 1  0  0 |   = | S[k]                      S[k-1] + S[k-2]    S[k-1] |
| 0  1  0 |     | S[k-1]                    S[k] - S[k-1]      S[k-2] |

And from this, by multiplying k-th and m-th powers of the matrix you can get the following relations:

S_k+m = S_kS_m-2 + S_k-2S_m + S_k-1S_m-2 + S_k-2S_m-1 + S_k-2S_m-2 + S_k-1S_m-1
S_k+m-1 = S_kS_m-1 + S_k-1S_m + S_k-1S_m-1 + S_k-2S_m-2
S_k+m-2 = S_k-1S_m-2 + S_k-2S_m-1 + S_kS_m – S_k-1S_m-1

Putting all of this together into code you get the following solution:

public class Solution {
    private static final int M = 1_000_000_007;

    // The set of equations for S[k+m], S[k+m-1], S[k+m-2]
    // a[t] is S[k-t] and b[t] is S[m-t]
    private static long[] mul(long[] a, long[] b) {
        long r0110 = a[0]*b[1] + a[1]*b[0];
        long r0011 = a[0]*b[0] + a[1]*b[1];
        return new long[] {
            (a[0]*b[2] + a[2]*b[0] + r0110 + r0011) % M,
            (a[1]*b[2] + a[2]*b[1] + r0011) % M,
            (r0110 + a[2]*b[2] - a[1]*b[1]) % M
        };
    }

    public static long countWays(int n) {
        long[] result = {1, 0, 0};
        long[] square = {1, 0, 0};
        // Exponentiation by squaring
        for (int i = n; i != 0; i /= 2) {
            if (i % 2 == 1) result = mul(result, square);
            square = mul(square, square);
        }
        return result[0];
    }
}