Liyang/collatz.hs

Lengths of Collatz sequences

module Main where

import Prelude hiding ( sequence_ )
import Control.Applicative
import Control.Monad ( ap )
import Control.Monad.State ( State )
import qualified Control.Monad.State as State
import Data.Map ( Map )
import qualified Data.Map as Map
import Data.Foldable
import Data.Traversable
import System.Environment
import Text.Printf

instance Applicative (State alpha) where
	pure = return
	(<*>) = ap

collatz :: Integer -> State (Map Integer Integer) Integer
collatz n = case n of
	1 -> return 0
	_ -> do
		mm <- State.gets (Map.lookup n)
		case mm of
			Just m -> return m
			Nothing -> do
				m <- succ `fmap` collatz
					(if even n then n `div` 2 else 3 * n + 1)
				State.modify (Map.insert n m)
				return m

main :: IO ()
main = getArgs >>= mainArgs

mainArgs :: [String] -> IO ()
mainArgs [a, b] | kosher ra && kosher rb && na <= nb && na > 0 = do
	let ns = [na .. nb]
	let ms = State.evalState (traverse collatz ns) Map.empty
	sequence_ (zipWith (printf "%d %d\n") ns ms) -- for Grapher.app
	-- sequence_ (zipWith (printf "%d takes %d iterations to reach 1\n") ns ms)
	where
	ra@ ~[(na, _)] = reads a
	rb@ ~[(nb, _)] = reads b
	kosher r = case r of
		[_] -> True
		_   -> False
mainArgs _ = return ()