配列版エラトステネスのふるい

sieve2.hs

import Control.Monad (mapM_, forM_, when)
import qualified Data.Array.IArray as IArray
import Data.Array.Unboxed (UArray)
import Data.Array.ST (runSTUArray)
import Data.Array.MArray (newArray, readArray, writeArray)

-- 与えられた数以下の素数
primes :: Int -> [Int]
primes = map fst . filter snd . IArray.assocs . sieve

-- i番目の要素にiが素数ならTrue、素数でなければFalseの入った0からnまでの配列を返す
sieve :: Int -> UArray Int Bool
sieve n = runSTUArray $ do
  table <- newArray (0, n) True
  mapM_ (\i -> writeArray table i False) $ [0, 1] ++ [4, 6 .. n]
  forM_ (takeWhile (\i -> i * i <= n) [2 ..]) $ \i -> do
    p <- readArray table i
    when p $
      forM_ [i * i, i * (i + 1) .. n] $ \j ->
        writeArray table j False
  return table

-- 200万以下の素数の総和を求める
main :: IO ()
main = print . sum . map (toInteger) . primes $ 2000000