Optimized sieve of Eratosthenes

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Array.Unboxed


soeST :: forall s. Int -> ST s (STUArray s Int Bool)
soeST n = do
    arr <- newArray (0, n) True
    mapM_ (\i -> writeArray arr i False) [0, 1]
    let n2 = n `div` 2

    let loop :: Int -> ST s ()
        loop i | i > n2 = return ()
        loop i = do
            b <- readArray arr i

            let reset :: Int -> ST s ()
                reset j | j > n = return ()
                reset j = writeArray arr j False >> reset (j + i)

            when b (reset (2*i))

            loop (succ i)

    loop 2
    return arr


soeA :: Int -> UArray Int Bool
soeA n = runST (soeST n >>= freeze)


soe :: Int -> [Int]
soe = map fst . filter snd . assocs . soeA


soeCount :: Int -> Int
soeCount = length . filter id . elems . soeA


main :: IO ()
main = print (soeCount (10^8))

Refinement of soeST

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
soeST :: forall s. Int -> ST s (STUArray s Int Bool)
soeST n = do
    arr <- newArray (0, n) True
    mapM_ (\i -> writeArray arr i False) [0, 1]
    let n2 = floor (sqrt (fromIntegral n))

    let loop :: Int -> ST s ()
        loop i | i > n2 = return ()
        loop i = do
            b <- readArray arr i

            let reset :: Int -> ST s ()
                reset j | j > n = return ()
                reset j = writeArray arr j False >> reset (j + i)

            when b (reset (i*i))

            loop (succ i)

    loop 2
    return arr