CarlEdman.hs

CarlEdman.hs

import qualified Control.Monad as ControlMonad
import qualified Data.Array as Array
import qualified Data.Array.ST as ArrayST
import qualified Data.Array.Base as ArrayBase
import           Control.DeepSeq
import qualified Data.Map as Map

-- From "The Genuine Sieve of Eratosthenes" by Melissa O'Neill.
-- O(n * log n * log log n)
sieve xs = sieve' xs Map.empty
           where sieve' []     table = []
                 sieve' (x:xs) table = case Map.lookup x table of
                                       Nothing    -> x : sieve' xs (Map.insert (x*x) [x] table)
                                       Just facts -> sieve' xs (foldl reinsert (Map.delete x table) facts)
                                                     where reinsert table prime = Map.insertWith (++) (x + prime) [prime] table

-- Stream of n prime numbers: 2, 3, 5, 7, 11, ...
-- O(n * log n * log log n)
primes = sieve [2..]

integerSquareRoot = floor . sqrt . fromIntegral

primeLucy :: (Integer -> Integer) -> (Integer -> Integer) -> Integer -> (Integer->Integer)
primeLucy f sf n = g
  where
    r = fromIntegral $ integerSquareRoot n
    ni = fromIntegral n
    loop from to c = let go i = ControlMonad.when (to<=i) (c i >> go (i-1)) in go from

    k = ArrayST.runSTArray $ do
      k <- ArrayST.newListArray (-r,r) $ force $
        [sf (div n (toInteger i)) - sf 1|i<-[r,r-1..1]] ++
        [0] ++
        [sf (toInteger i) - sf 1|i<-[1..r]]
      ControlMonad.forM_ (takeWhile (<=r) primes) $ \p -> do
        l <- ArrayST.readArray k (p-1)
        let q = force $ f (toInteger p)

        let adjust = \i j -> do { v <- ArrayBase.unsafeRead k (i+r); w <- ArrayBase.unsafeRead k (j+r); ArrayBase.unsafeWrite k (i+r) $!! v+q*(l-w) }

        loop (-1)         (-div r p)              $ \i -> adjust i (i*p)
        loop (-div r p-1) (-min r (div ni (p*p))) $ \i -> adjust i (div (-ni) (i*p))
        loop r            (p*p)                   $ \i -> adjust i (div i p)

      return k

    g :: Integer -> Integer
    g m
      | m >= 1 && m <= integerSquareRoot n                       = k Array.! (fromIntegral m)
      | m >= integerSquareRoot n && m <= n && div n (div n m)==m = k Array.! (fromIntegral (negate (div n m)))
      | otherwise = error $ "Function not precalculated for value " ++ show m

primeSum :: Integer -> Integer
primeSum n = (primeLucy id (\m -> div (m*m+m) 2) n) n

main = print $ primeSum (2*10^9)