23 February 2012

Calculating Combinations of 1, 2, 3 Step Jumps

So again I was trolling along on Antonio Gulli's website (which has become a bit of a favorite to find fun problems to solve during a quiet evening). I came across this question:
A child can hop either 1, 2, or 3 steps in a stairway of N steps.
How many combinations do you have?
As the first comment mentions then this problem can be solved using recursion. But the question is, how would you translate this into program code?

... want a second to think about this before reading my attempts?

Initial Attempt
After thinking about this problem for a little while I decided that this might be cleanly solved using a recursive tree traversal algorithm. As we have a branching factor of 3 we should not have a branch height of more than N/3 (only taking 3 step jumps) in the best case and N in the worst case (if only taking 1 step jumps).

The initial naive solution I came up with was a simple DFS search using a stack (LIFO):

static Stack<long> nodes = new Stack<long>();
static long Combinations( int n )
{
    long leafs = 0;
    nodes.Push(n);
    while (nodes.Count > 0)
    {
        var x = nodes.Pop();

        // Leaf, just count
        if (x <= 0) leafs++;
        else
        {
            // Only push down options where the child 
            // can actually complete the stair jump
            if (x - 1 >= 0) nodes.Push(x - 1);
            if (x - 2 >= 0) nodes.Push(x - 2);
            if (x - 3 >= 0) nodes.Push(x - 3);
        }
    }
    return leafs;
}

Improved Attempt
It is soon obvious if you run the first code using non-trivial sizes of N that this naive DFS solution is simply too slow. The time that it takes to iterate through the entire tree and generating all the intermediary nodes is just too great and quite quickly the algorithm takes too long to return a solution within a reasonable time.

When you look closer at this ternary tree it might become obvious that since it is a complete tree it has a simple repeating pattern down each of its branches. Basically each branch is a sub-tree of the one preceding it (with node count n-1). Meaning that the (n-1) branch far to the left at height h has an equal amount of nodes in its immediate (n-1) child at h-1 as the second (n-2) branch at h has.

Using this we can simply pre-cache each height in the tree once we have calculated it completely and thus really only need to do one full depth run on a single branch after which we will have all the necessary values calculated for the remaining nodes.

           A
          /|\
         B C D
        /|\
       E F G
      /|\
     H I J
We only need to calculate the A,B,E,H values, after that no extra recursive traversal is necessary into any other subtrees as we can use the cached value for H for F, J for G, E for C, F for D and so forth.

So by adding a simple map that caches the node count for every remaining node number we can significantly speed up this calculation:

static Dictionary<int, BigInteger> map = new Dictionary<int, BigInteger>();
static BigInteger Combinations2(int n)
{
    if (n <= 0) return 1;

    BigInteger leafs = 0;
    if (map.TryGetValue(n, out leafs))
        return leafs;

    // We need to calculate it if not found
    if (n - 1 >= 0)
        leafs += Combinations2(n - 1);
    if (n - 2 >= 0)
        leafs += Combinations2(n - 2);
    if (n - 3 >= 0)
        leafs += Combinations2(n - 3);

    map[n] = leafs;
    return leafs;
}

We can probably do better than this though?