Memoized solution to Count number of ways to climb n steps

Question

Problem statement : A Child is climbing up n steps with either 1 , 2, 3 hops ..how many ways can the child climb up the stairs?

source: cracking the coding interview book.

I have two solutions the brute force and the memoized one.

Wanted to know how to improve my solution further, any issues you see with this code?

description of the code: Brute force: basically count in a recursive manner, as there are 3 ways to reach the nth step so we count number of ways to reach each of n-1 , n-2 and n-3 steps and then sum them up to find the number of ways of reaching the nth step.

memoized solution:

we are now caching the solutions in a hashmap to reuse the solutions.

The Code:

public class CountNWaysforSteps {

    // running time O(3 ^ n) inefficient solution 
    public static int numberOfWays(int n){

        if(n < 0 ) return 0;
        if(n == 0 ) return 1;

        return numberOfWays(n-1) + numberOfWays(n-2) + numberOfWays(n-3);
    }
    //memoized solution running time O(N)
    public static int numberOfWaysMemoized(int n, Map<Integer,Integer> cache){

        if(n < 0 ) return 0;
        if(n == 0 ) return 1;

        if(cache.containsKey(n)){
            return cache.get(n);
        }

        int calculatedNumberOFWays = numberOfWaysMemoized(n-1, cache) + numberOfWaysMemoized(n-2, cache) + numberOfWaysMemoized(n-3, cache);
        cache.put(n,calculatedNumberOFWays);
        return calculatedNumberOFWays;
    }

    public static void main(String...args){
        System.out.println(numberOfWays(5));
        System.out.println(numberOfWaysMemoized(5, new HashMap<>()));
    }
}

Answer 1

First impressions:

1. The dynamic approach is a good idea, but it can be improved
2. When a Map<Integer, T> is used, see if a T[] can be used instead
3. When doing a dynamic solution, see if you can remove the need for recursion

By applying these tricks, my suggested solution is this:

public static int numberOfWaysFast(int n) {
    if (n < 2) {
        return n >= 0 ? 1 : 0;
    }
    int[] cache = new int[n+1];
    cache[0] = 1;
    cache[1] = 1;
    cache[2] = 2;
    for (int i = 3; i <= n; i++) {
        cache[i] = cache[i-1] + cache[i-2] + cache[i-3];
    }
    return cache[n];
}

This solution is about 50-100 times faster on my computer, and doesn't run into StackOverflowErrors for large inputs.

If you want to improve even further, you can reuse the cache array for subsequent calls, and only calculate new values when you have to. This could be applicable if you knew the maximum bound for values to be calculated, and had to rerun the function multiple times. Otherwise you could use an ArrayList<Integer>, but that's a tiny bit slower.

EDIT: forgot to handle input <2, added some code for it.

Answer 2

Caller friendliness:

numberOfWaysMemoized() puts a burden on the caller of passing in the required cache for the function.

This can be improved by providing a second function which creates the cache for the caller:

public static int numberOfWaysMemoized(int n) {
    return numberOfWaysMemoized(n, new HashMap<>());
}

private static int numberOfWaysMemoized(int n, Map<Integer,Integer> cache) {
    //...
}

Alternately, provide a global cache, so subsequent callers can get the benefit of already cached results:

private static Map<Integer, Integer> cache = new HashMap<>();

public static int numberOfWaysMemoized(int n) {
    //...
}

---

Double lookups in Map<Integer,Integer>:

if(cache.containsKey(n)){
    return cache.get(n);
}

First, n is looked up in cache, just to see if it exists. Then, if it was found, it is looked up a second time to fetch the value.

This can be combined into just one lookup:

Integer value = cache.get(n);
if (value != null)
    return value;

If the value n was not found, we still lookup a second time where it would go in cache as a side effect of storing the value:

cache.put(n,calculatedNumberOFWays);

The containsKey(), get() and put() can all be reduced to just one cache lookup operation, using computeIfAbsent(). For example:

return cache.computeIfAbsent(n, n -> numberOfWaysMemoized(n-1, cache)
                                   + numberOfWaysMemoized(n-2, cache)
                                   + numberOfWaysMemoized(n-3, cache));

---

Improving on @maxb's solution:

That solution is O(n) complexity in time, but is also O(n) complexity in space. It allocates storage which is only briefly used; once cache[i] is computed, cache[i-3] is never used again.

The entire cache may be maintained in 3 local variables, yielding O(1) memory usage:

public static int numberOfWaysFaster(int n) {
    if (n < 2) {
        return n >= 0 ? 1 : 0;
    }
    int cache0 = 1;
    int cache1 = 1;
    int cache2 = 2;
    for (int i = 3; i <= n; i++) {
        int tmp = cache0 + cache1 + cache2;
        cache0 = cache1;
        cache1 = cache2;
        cache2 = tmp;
    }
    return cache2;
}

The ugly part of the above solution is the memory shifting. The cache values all have to shift down by one location, on each iteration. We can avoid this by loop-unrolling:

int cache3 = cache0 + cache1 + cache2;
int cache4 = cache1 + cache2 + cache3;
int cache5 = cache2 + cache3 + cache4;

Note that cache0 is last used to compute cache3, and then is no longer used. Also, cache1 is last used to compute cache4, and cache2 is last used to compute cache5. This means we can use the variable cache0 to store the value of cache3, cache1 to store cache4, and cache2 to store cache5. No memory shifting is required:

public static int numberOfWaysFastest(int n) {
    int cache0 = 1;
    int cache1 = 1;
    int cache2 = 2;

    while (n >= 3) {
        cache0 = cache0 + cache1 + cache2;
        cache1 = cache1 + cache2 + cache0;
        cache2 = cache2 + cache0 + cache1;
        n -= 3;
    }

    switch (n) {
        case 0: return cache0;
        case 1: return cache1;
        case 2: return cache2;
    }

    return 0;
}