How do I write a generic memoize function?

Question

I'm writing a function to find triangle numbers and the natural way to write it is recursively:

function triangle (x)
   if x == 0 then return 0 end
   return x+triangle(x-1)
end

But attempting to calculate the first 100,000 triangle numbers fails with a stack overflow after a while. This is an ideal function to memoize, but I want a solution that will memoize any function I pass to it.

Answer 1

In the vein of posting memoization in different languages, i'd like to respond to @onebyone.livejournal.com with a non-language-changing C++ example.

First, a memoizer for single arg functions:

template <class Result, class Arg, class ResultStore = std::map<Arg, Result> >
class memoizer1{
public:
    template <class F>
    const Result& operator()(F f, const Arg& a){
        typename ResultStore::const_iterator it = memo_.find(a);
        if(it == memo_.end()) {
            it = memo_.insert(make_pair(a, f(a))).first;
        }
        return it->second;
    }
private:
    ResultStore memo_;
};

Just create an instance of the memoizer, feed it your function and argument. Just make sure not to share the same memo between two different functions (but you can share it between different implementations of the same function).

Next, a driver functon, and an implementation. only the driver function need be public int fib(int); // driver int fib_(int); // implementation

Implemented:

int fib_(int n){
    ++total_ops;
    if(n == 0 || n == 1) 
        return 1;
    else
        return fib(n-1) + fib(n-2);
}

And the driver, to memoize

int fib(int n) {
    static memoizer1<int,int> memo;
    return memo(fib_, n);
}

Permalink showing output on codepad.org. Number of calls is measured to verify correctness. (insert unit test here...)

This only memoizes one input functions. Generalizing for multiple args or varying arguments left as an exercise for the reader.

Answer 2

Here's something that works without converting the arguments to strings. The only caveat is that it can't handle a nil argument. But the accepted solution can't distinguish the value nil from the string "nil", so that's probably OK.

local function m(f)
  local t = { }
  local function mf(x, ...) -- memoized f
    assert(x ~= nil, 'nil passed to memoized function')
    if select('#', ...) > 0 then
      t[x] = t[x] or m(function(...) return f(x, ...) end)
      return t[x](...)
    else
      t[x] = t[x] or f(x)
      assert(t[x] ~= nil, 'memoized function returns nil')
      return t[x]
    end
  end
  return mf
end

Answer 3

function memoize (f)
   local cache = {}
   return function (x)
             if cache[x] then
                return cache[x]
             else
                local y = f(x)
                cache[x] = y
                return y
             end
          end
end

triangle = memoize(triangle);

Note that to avoid a stack overflow, triangle would still need to be seeded.

Answer 4

Update: Commenters have pointed out that memoization is a good way to optimize recursion. Admittedly, I hadn't considered this before, since I generally work in a language (C#) where generalized memoization isn't so trivial to build. Take the post below with that grain of salt in mind.

I think Luke likely has the most appropriate solution to this problem, but memoization is not generally the solution to any issue of stack overflow.

Stack overflow usually is caused by recursion going deeper than the platform can handle. Languages sometimes support "tail recursion", which re-uses the context of the current call, rather than creating a new context for the recursive call. But a lot of mainstream languages/platforms don't support this. C# has no inherent support for tail-recursion, for example. The 64-bit version of the .NET JITter can apply it as an optimization at the IL level, which is all but useless if you need to support 32-bit platforms.

If your language doesn't support tail recursion, your best option for avoiding stack overflows is either to convert to an explicit loop (much less elegant, but sometimes necessary), or find a non-iterative algorithm such as Luke provided for this problem.

Answer 5

Mathematica has a particularly slick way to do memoization, relying on the fact that hashes and function calls use the same syntax:

triangle[0] = 0;
triangle[x_] := triangle[x] = x + triangle[x-1]

That's it. It works because the rules for pattern-matching function calls are such that it always uses a more specific definition before a more general definition.

Of course, as has been pointed out, this example has a closed-form solution: triangle[x_] := x*(x+1)/2. Fibonacci numbers are the classic example of how adding memoization gives a drastic speedup:

fib[0] = 1;
fib[1] = 1;
fib[n_] := fib[n] = fib[n-1] + fib[n-2]

Although that too has a closed-form equivalent, albeit messier: http://mathworld.wolfram.com/FibonacciNumber.html

I disagree with the person who suggested this was inappropriate for memoization because you could "just use a loop". The point of memoization is that any repeat function calls are O(1) time. That's a lot better than O(n). In fact, you could even concoct a scenario where the memoized implementation has better performance than the closed-form implementation!

Answer 6

I bet something like this should work with variable argument lists in Lua:

local function varg_tostring(...)
    local s = select(1, ...)
    for n = 2, select('#', ...) do
        s = s..","..select(n,...)
    end
    return s
end

local function memoize(f)
    local cache = {}
    return function (...)
        local al = varg_tostring(...)
        if cache[al] then
            return cache[al]
        else
            local y = f(...)
            cache[al] = y
            return y
        end
    end
end

You could probably also do something clever with a metatables with __tostring so that the argument list could just be converted with a tostring(). Oh the possibilities.