HackerRank Recursion: Davis' Staircase

The point of this HackerRank problem is calculating a sort of uber-Fibonacci number. Since we want to have an efficient solution, we should immediately think to a dynamic programming approach, or at least to some kind of memoization.

Let's call Davis number the number of ways we can climb a staircase, with the rule that we can take 1, 2, or 3 steps at a time.

First thing, we want to identify the base cases. That's pretty easy.

1. Our staircase has just one step, we have only one way to climb it, taking 1 step.
2. We can climb it in a single two-step jump, or we can take 1 step, then falling back to the previous case. Total of two ways of climbing.
3. A single three-step jump is a way to get on top, otherwise we can take one step and falling back to case (2) or take a two-step jump and fall back to case (1). This means 1 + 2 + 1 = 4 ways of climbing this staircase.

The generic case is determined by the three precedent ones, something like this:

result = davis_number(n-1) + davis_number(n-2) + davis_number(n-3)

Given that, I skipped the step of writing a plain recursive solution to the problem, and I moved directly to write a Python script that uses a cache to get the result via DP.

I cached the basic step results as identified above in a list (and I define a constant that make explicit the requirement stating that a staircase could not have more that 36 steps):

cache = [1, 2, 4]
MAX_STEPS = 36

And then I wrote my solution:

def solution(steps):
    assert 0 < steps <= MAX_STEPS  # 1
    if steps <= len(cache):  # 2
        return cache[steps - 1]

    for _ in range(steps - len(cache)):  #3
        cache.append(cache[-1] + cache[-2] + cache[-3])

    return cache[-1]  # 4

1. Without this assertion, the code would become unstable for negative input values (maybe an IndexError exception, maybe simply a wrong result) and would take a very long time for large positive input values. Better a predictable AssertionError instead.
2. The result has been already calculated, and it is available in the cache. Just return it.
3. Let's calculate the required value. To do that, we need to calculate all the previous values, so we extend the cache calculating in order all of them.
4. Finally, we return the rightmost element in the cache, that is, the required Davis' number.

If MAX_STEPS was a bigger number, it could be a good idea to predetermine the full size for the cache, filling it with zeroes, and pushing the actual values when the solutions are created. This would require an extra variable to keep track of the current last Davis number in cache, and would make the code slightly more complicated. Here I found better to keep the code simple.

Full Python script and testcase on GitHub.