Staircase problem solved using recursion

Question

Problem

Question taken from here.

I want to go up a flight of stairs that has n steps. I can either take 1 or 2 steps each time. How many different ways can I go up this flight of stairs? Write a function count_stair_ways that solves this problem for me. def count_stair_ways(n):

Solution:

def count_stair_ways(n):
    if n == 1: # 1 way to flight 1 stair case step
        return 1
    if n ==2: # 2 ways to flight 2 stair case steps(1+1, 2)
        return 2
    return count_stair_ways(n-1) + count_stair_ways(n-2)

Can we improve this code?

Answer 1

The staircase problem actually just generates the Fibonnacci sequence.

Whilst the recursive solution is nice, without memoization you're much better off just using a loop:

def count_stairways(n):

   a, b = 0, 1
   for _ in range(n):
       a, b = b, a+b

   return b

A nice alternative if you want multiple values out is to create a generator:

def count_stairways():

    a, b = 0, 1
    while True:
        a, b = b, a+b
        yield b

If you just want a particular value (possibly for a large n), the fibonnacci sequence actually has a closed form

 def count_stairways(n):

     phi = (1 + math.sqrt(5)) / 2

     return int((pow(phi, n+1) - pow(1-phi, n+1))/math.sqrt(5))

Answer 2

Your code is easy to understand and just works but it can still be improved.

Documentation

It could be nice to add some documentation to tell which values are expected.

Domain

When I tried with 0 as an input, I got a RuntimeError: maximum recursion depth exceeded. This is because the recursion calls go deeper and deeper into negative numbers without reaching a stop.

Adding :

if n < 0:
    return 0
if n == 0:
    return 1

just solves this problem. By the way, there is 1 way to flight a 0 stair case step.

Tests

It is a good habit to add tests for your code. I won't dive into testing framework and everything but just adding :

if __name__ == "__main__":
    for n, value in enumerate([1, 1 ,2 ,3 ,5 ,8 ,13 ,21 ,34 ,55]):
        res = count_stair_ways(n)
        assert res == value, "exp:%d got:%d for n = %d" % (value, res, n)

gives you a nice way to see if things ever go wrong.

Keep it simple

Once you handle all n <= 0, you don't need to handle the cases 1 and 2 individually. You can remove :

if n == 1: # 1 way to flight 1 stair case step
    return 1
if n ==2: # 2 ways to flight 2 stair case steps(1+1, 2)
    return 2

and see that it still works fine.

At this point, the code becomes :

def count_stair_ways(n):
    """ Some doc. Input can be any integer n."""
    if n < 0:
        return 0
    if n == 0:
        return 1
    return count_stair_ways(n-1) + count_stair_ways(n-2)

if __name__ == "__main__":
    for n, value in enumerate([1, 1 ,2 ,3 ,5 ,8 ,13 ,21 ,34 ,55]):
        res = count_stair_ways(n)
        assert res == value, "exp:%d got:%d for n = %d" % (value, res, n)

Performance

Your solution will be slow for any big (or even medium) inputs. You can easily see add a print statement at the beginning of the function to see how many times it gets called and you'll see that we compute the same things many times.

When n is 15, the function gets called that many times with the different inputs :

610 -1
987 0
610 1
377 2
233 3
144 4
 89 5
 55 6
 34 7
 21 8
 13 9
  8 10
  5 11
  3 12
  2 13
  1 14
  1 15

There must be a better way. Using memoization is the way to go but there's an even easier option involving a simple loop.

def count_stair_ways(n):
    """ Some doc. Input can be any integer n."""
    a, b = 0, 1
    for _ in range(n + 1):
        a, b = b, a + b
    return a

Going further

You sequence actually corresponds to Fibonacci numbers. You'll find various super efficient way to compute them.

Answer 3

This solution takes a problem of size n and decomposes it into two problems of approximately size n - 1. The complexity would therefore be O(2ⁿ). Furthermore, many of those computations would be performed over and over again, so the solution would benefit greatly from memoization.

Answer 4

To get style out of the way:

You should have a space after == in if n ==2.

---

You can combine the first two if-else statements into one statement with:

if n in (1, 2):
    return n

Like @200_success said, you should be caching your results, so that way, when you keeping increasing, you won't have to recalculate every thing over and over.