A bit of Mathematica

Today’s post was motivated by a problem at Project Euler involving the Collatz conjecture. Since this is a little bit spoiler-y, I’ve included the main body of the post after the cut. Don’t read it unless you’ve either solved the problem at Project Euler or you’re actually looking for spoilers. You have been warned.

Recently, I solved Problem 14 at Project Euler using a few lines of Mathematica code. Unfortunately, I was disappointed in the performance of my solution, since the problems are designed to be able to be solved by computer in less than 1 minute. My initial code took around 8.5 minutes, but it did generate the correct answer.

So, I set out to see if I could improve its performance.

Here was my initial implementation of the problem:

Collatz[1] = 1;
Collatz[n_?OddQ] := 3 n + 1;
Collatz[n_?EvenQ] := n/2;

CollatzSequence[n_] := Most[FixedPointList[Collatz, n]];

CollatzLength[n_] := Length[CollatzSequence[n]];

Problem14[n_] := Block[{},
   results = Map[{#, CollatzLength[#]} &, Range[n]];
   m = Max[Map[First[Rest[#]] &, results]];
   Return  [Map [First, Select[results, #[[2]] == m &]][[1]]]];

However, when I did

Timing [Problem14[10^6]]

the results were less than encouraging: I got the correct result, but the code took a little over 9 minutes to return on my 1.7 ghz dual-core Athlon laptop. Since Project Euler problems are designed to be solved using around a minute’s worth of computing time, I wasn’t too happy about that.

What I did was replace the definition of CollatzLength with the following definition:

Block [{cache},
    CollatzLength[1] := 1;
    cache[1] = 1;
    CollatzLength[n_] := Block [{$RecursionLimit = Infinity},
        If [Not[NumberQ [cache[n]]],
	cache[n] = CollatzLength[Collatz[n]] + 1;cache[n],
	cache[n]]];
];

Using this function instead of the naïve function allowed me to complete the calculation in 1 minute 12 seconds — a factor of 7.5 improvement. Not too shabby, if I do say so myself.

What I did, in a nutshell, was replace the original function with a recursive function that caches previously calculated values. I suspect this might have actually been faster if I had implemented it in a compiled Scheme dialect, but, at this point, since the code was running in approximately 1 minute, I was more or less happy with it.

The tradeoff, of course, was that the faster function uses approximately 6 times the memory of the original function. However, when the difference is between 50 megs and 300 megs, on a system with 2 gigs of memory, it seems like a good tradeoff to me.