Play With Lua!

Or, how to annihilate an interview question. I was asked this a little over a year ago in an interview, and I’d been asked it a few times before then too; it’s pretty common.

“Write a function that takes a parameter n and returns the nth Fibonacci number.” The Fibonacci sequence, if you don’t already know, is defined like this: f(1) and f(2) are both 1. After that, each number in the sequence is the sum of the previous two, so, f(3) is 2 (because 1+1), then f(4) is 3, then f(5) is 5, and so on.

(Fun, useless fact: this was originally discovered as part of a thought experiment about rabbits. If you imagine that you start with two rabbits, and each month each rabbit produces another rabbit (but rabbits don’t reproduce until they’re one month old), then the rabbit population after n months is the nth Fibonacci number.)

Naive version

So, start off by writing the naive recursive version:

function fib(n)
    if n < 3 then return 1
    else return fib(n-1) + fib(n-2) end
end

This works, but it’s inefficient: if you call it with n=4, say, it results in 5 calls: 4, 3, 2, 1, and 2. Calling it with n=8 is 21 calls. In fact, it’s actually O(fib(n)) complexity; the number of recursive calls rises at the same rate as the Fibonacci sequence itself.

So let’s make it faster.

Iterative version

One way is to make it iterative:

function ifib(n)
    local a,b = 1,1
    for k = 3,n do
        a, b = b, a+b
    end
    return b
end

Now it won’t recurse at all, of course, but it’s kind of ugly. Recursion is cool, iteration is boring. Let’s try to make the recursive one more efficient.

Memoized version

If you look at the naive version, it’s calling itself with the same argument several times. We can help matters by caching the results so after it calls once it never has to again:

function memoize(fn)
    local cache = {}
 
    return function(n)
        cache[n] = cache[n] or fn(n)
        return cache[n]
    end
end
 
fib = memoize(fib)

Now when you call it with n=4, say, you only get 4 recursive calls: 4, 3, 2, and 1. In general it will call itself n times, because if it ever has to repeat a call it’ll use the cache. So now it’s O(n).

But that’s not perfect. We’re still using a lot of stack space. It would be nicer if it ran in constant stack space.

Tail-recursive version

function tfib(n)
    local function f(a, b, n)
        if n < 3 then
            return b
        else
            return f(b, a+b, n-1)
        end
    end
 
    return f(1,1,n)
end

You can see this is a lot like the iterative version: we’re starting from the first two values and working up, and shifting (a+b)-to-b and b-to-a each time. But the mechanism of looping is recursion, and the recursive call is the last thing that happens in any given invocation, so, it’s tail recursive and Lua can eliminate the tail call so it never blows out the stack.

At this point we’ve got a naive, easy-to-read version; a version that’s fast and small but ugly; a version that’s fast and pretty but needs memory for a cache; and a version that’s fast and pretty and uses constant memory. The only way to improve it further is to be clever:

The closed-form version

The first three new versions we made are general transformations useful for speeding up any slow function. This one is specific to the Fibonacci function. It turns out that as n approaches infinity, the ratio of fib(n) to fib(n-1) is phi, the golden ratio.

Phi is (1 + math.sqrt(5)) / 2. The closest integer to phi^n / math.sqrt(5) is the nth Fibonacci number, and we can easily find that with math.floor:

phi = (1+math.sqrt(5)) / 2
 
function fib(n)
    return math.floor(phi^n / math.sqrt(5) + 1/2)
end

It doesn’t get much better than this: constant time and constant space. Beyond this there’s not a lot you can do to make it more efficient.