Skip to content

Instantly share code, notes, and snippets.

@AdamStelmaszczyk
Created June 9, 2017 12:42
Show Gist options
  • Save AdamStelmaszczyk/e0e794e6392dd7ff9758fc64d5557622 to your computer and use it in GitHub Desktop.
Save AdamStelmaszczyk/e0e794e6392dd7ff9758fc64d5557622 to your computer and use it in GitHub Desktop.
After Shersh code review
import Data.List
import Data.IntMap.Strict(IntMap, (!))
import qualified Data.IntMap.Strict as IntMap
problem10 :: Int -> Int
problem10 n = (sieve (IntMap.fromList [(i, i * (i + 1) `div` 2 - 1) | i <- vs]) 2 r vs) ! n
where vs = [n `div` i | i <- [1..r]] ++ [n', n' - 1 .. 1]
r = floor (sqrt (fromIntegral n))
n' = n `div` r - 1
sieve :: IntMap Int -> Int -> Int -> [Int] -> IntMap Int
sieve m p r vs | p > r = m
| m ! p > m ! (p - 1) = sieve (update m vs p) (p + 1) r vs
| otherwise = sieve m (p + 1) r vs
update :: IntMap Int -> [Int] -> Int -> IntMap Int
update m vs p = foldl' decrease m (takeWhile (>= p*p) vs)
where sumOfSieved v = p * (m ! (v `div` p) - m ! (p - 1))
decrease m v = IntMap.adjust (subtract $ sumOfSieved v) v m
main = print $ problem10 $ 2*10^9
-- $ time ./10
-- 95673602693282040
-- real 0m3.279s
-- user 0m3.231s
-- sys 0m0.041s
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment