problem14.hs

-- http://projecteuler.net/index.php?section=problems&id=14
{-# OPTIONS_GHC -O2 #-}
import Data.List (maximumBy)
import Data.Ord (comparing)
import Control.Applicative ((<$>))
import Control.Monad (forM_)
import Data.Array.Unboxed (UArray, assocs)
import Data.Array.ST (runSTUArray, newArray, readArray, writeArray)

collatz :: Integer -> Integer
collatz n
  | even n = n `div` 2
  | otherwise = 3*n + 1

collatzLens :: Int -> UArray Int Int
collatzLens n = runSTUArray $ do
  memo <- newArray (0,n) 0
  writeArray memo 1 1
  forM_ [2 .. toInteger n] $
    let go i
          | i > toInteger n = (+1) <$> go (collatz i)
          | otherwise = do
              v <- readArray memo (fromIntegral i)
              if v /= 0
                then return v
                else do
                  w <- (+1) <$> go (collatz i)
                  writeArray memo (fromIntegral i) w
                  return w
    in go
  return memo

main = print . solve $ 1000000
  where
    solve = fst . maximumBy (comparing snd) . assocs . collatzLens