memorize.hs

import qualified Data.Map as Map
import Control.Monad.State (state, runState)
import qualified Control.Monad.State as State
import Control.Monad (foldM, ( <=< ))

collatz :: ((Int -> a) -> Integer -> a) -> (Int -> a) -> Integer -> a
collatz f cont n
  | n == 1 = cont 1
  | even n = f (cont . succ) $ n `div` 2
  | odd n  = f (cont . succ) $ 3 * n + 1

-- 型が長すぎて書く気がしますん
memorize f n = memorize' f (memo n) n
  where
    memo n m = state $ \cache -> (m, Map.insert n m cache)
    memorize' f cont n = do
      cache <- State.get
      case Map.lookup n cache of
        Just m -> cont m
        Nothing -> f (memorize' f) (cont <=< memo n) n

main :: IO ()
main = print . fst . flip runState Map.empty . foldM (\n n' -> do
  m <- memorize collatz n
  m' <- memorize collatz n'
  return $! if m > m' then n else n') 1 $ [2 .. 999999]