Project Euler 78

http://projecteuler.net/index.php?section=problems&id=78

まともに計算すると遅いので、メモ化する。
コンパイルするときは --makeをつける。それでもPython並みに遅い。

import Data.Map (Map, (!), empty, insert, lookup)

n = 10^6

f m 0 = (insert 0 1 m, 1)
f m n = case Data.Map.lookup n m of
        Just p  -> (m, p)
        Nothing -> (insert n q m, q) where
                q = (g m 1 (\k -> n - (div (k * (3 * k - 1)) 2)))
                  + (g m 1 (\k -> n - (div (k * (3 * k + 1)) 2))) where
                        g m k h | h k < 0   = 0
                                | otherwise = (m!(h k)) - (g m (k + 1) h)

calc m k | mod fk n == 0 = (m1, k)
         | otherwise     = calc m1 (k + 1)
                                where (m1, fk) = f m k

main = print (snd (calc empty 0))