Project Euler 95

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

次を配列にしておく。

import qualified Data.Map as M
import qualified Data.Set as S
import Data.Array
import Data.List

sieve max_n = update 2 (M.fromList [ (n,n) | n <- [2..max_n] ]) where
        update n m | n * n > max_n  = m
                   | (m M.! n) == n = update (n + 1) (f n (n * 2) m)
                   | otherwise      = update (n + 1) m where
                f p q m | q > max_n = m
                        | m M.! q == q  = f p (q + p) (M.insert q p m)
                        | otherwise = f p (q + p) m

calc_exp n p | mod n p /= 0 = (n, 0)
             | otherwise    = (m, e + 1)
                        where (m, e) = calc_exp (div n p) p
factorize 1 = []
factorize n = [(p, e)] ++ (factorize q) where
                (q, e) = calc_exp n p
                p = m M.! n

next n = (product [ div (p^(e+1)-1) (p-1) | (p,e) <- factorize n ]) - n

chain_length n = f n n S.empty where
        f n m s | q == n       = (S.size s + 1)
                | q < n        = 0
                | q > max_n    = 0
                | S.member q s = 0
                | otherwise    = f n q (S.insert q s) where
                        q = next m

max_n = 10^6
m = sieve max_n
a = array (1,max_n) [ (n,next n) | n <- [1..max_n] ]
longer x y = if snd x >= snd y then x else y
main = print (fst (foldl' longer (0,0)
            [ (n,chain_length n) | n <- [1..max_n] ]))