algorithm - 如何优化这个 Haskell 代码在亚线性时间内总结素数？

Question

问题 10 来自 Project Euler就是在给定的 n 下找到所有素数的总和。

我简单地通过总结埃拉托色尼筛法生成的素数就解决了这个问题。然后我遇到了更有效率的solution by Lucy_Hedgehog (亚线性!)。

对于 n = 2⋅10^9:

• Python code (来自上面的引述)在 Python 2.7.3 中运行时间为 1.2 秒。

• C++ code (我的)运行时间约为 0.3 秒(使用 g++ 4.8.4 编译)。

我在 Haskell 中重新实现了相同的算法，因为我正在学习它:

import Data.List

import Data.Map (Map, (!))
import qualified Data.Map as Map

problem10 :: Integer -> Integer
problem10 n = (sieve (Map.fromList [(i, i * (i + 1) `div` 2 - 1) | i <- vs]) 2 r vs) ! n
              where vs = [n `div` i | i <- [1..r]] ++ reverse [1..n `div` r - 1]
                    r  = floor (sqrt (fromIntegral n))

sieve :: Map Integer Integer -> Integer -> Integer -> [Integer] -> Map Integer Integer
sieve m p r vs | p > r     = m
               | otherwise = sieve (if m ! p > m ! (p - 1) then update m vs p else m) (p + 1) r vs

update :: Map Integer Integer -> [Integer] -> Integer -> Map Integer Integer
update m vs p = foldl' decrease m (map (\v -> (v, sumOfSieved m v p)) (takeWhile (>= p*p) vs))

decrease :: Map Integer Integer -> (Integer, Integer) -> Map Integer Integer
decrease m (k, v) = Map.insertWith (flip (-)) k v m

sumOfSieved :: Map Integer Integer -> Integer -> Integer -> Integer
sumOfSieved m v p = p * (m ! (v `div` p) - m ! (p - 1))

main = print $ problem10 $ 2*10^9

我用 ghc -O2 10.hs 编译它并用 time ./10 运行。

它给出了正确答案，但大约需要 7 秒。

我用 ghc -prof -fprof-auto -rtsopts 10 编译它并用 ./10 +RTS -p -h 运行。

10.prof表明 decrease 需要 52.2% 的时间和 67.5% 的分配。

运行 hp2ps 10.hp 后，我得到了这样的堆配置文件:

再次看起来 decrease 占用了大部分堆。 GHC 版本 7.6.3。

您将如何优化此 Haskell 代码的运行时间？

---

更新 13.06.17:

我tried用 hashtables 包中的可变 Data.HashTable.IO.BasicHashTable 替换不可变的 Data.Map，但我可能做错了什么，因为对于微小的 n = 30，它已经花费了太长时间，大约 10 秒。怎么了？

更新 18.06.17:

Curious about the HashTable performance issues是一本好书。我拿了 Sherh 的 code使用可变 Data.HashTable.ST.Linear，但是 dropped Data.Judy in instead .它在 1.1 秒内运行，仍然相对较慢。

Answer 1

我做了一些小的改进，所以它在我的机器上运行 3.4-3.5 秒。 使用 IntMap.Strict 帮助很大。除此之外，我只是手动执行了一些 ghc 优化以确保安全。并通过您的链接使 Haskell 代码更接近 Python 代码。作为下一步，您可以尝试使用一些可变的 HashMap。但我不确定... IntMap 不会比某些可变容器快多少，因为它是一个不可变容器。尽管我仍然对它的效率感到惊讶。我希望这可以更快地实现。

代码如下:

import Data.List (foldl')
import Data.IntMap.Strict (IntMap, (!))
import qualified Data.IntMap.Strict as IntMap

p :: Int -> Int
p 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) :: Double)
                     n' = n `div` r - 1

sieve :: IntMap Int -> Int -> Int -> [Int] -> IntMap Int
sieve m' p' r vs = go m' p'
  where
    go m p | p > r               = m
           | m ! p > m ! (p - 1) = go (update m vs p) (p + 1)
           | otherwise           = go m (p + 1)

update :: IntMap Int -> [Int] -> Int -> IntMap Int
update s vs p = foldl' decrease s (takeWhile (>= p2) vs)
  where
    sp = s ! (p - 1)
    p2 = p * p
    sumOfSieved v = p * (s ! (v `div` p) - sp)
    decrease m  v = IntMap.adjust (subtract $ sumOfSieved v) v m

main :: IO ()
main = print $ p $ 2*10^(9 :: Int) 

更新:

使用可变的哈希表 我已经设法在 Haskell 上使用 this implementation 将性能提高到 ~5.5sec .

另外，我在几个地方使用了未装箱的向量而不是列表。 Linear 哈希似乎是最快的。我认为这可以更快地完成。我注意到 sse42 option在 hasthables 包中。不确定我是否已设法正确设置它，但即使没有它也运行得那么快。

更新 2(2017 年 6 月 19 日)

通过完全删除 judy hashmap，我设法使其 3x 比 @Krom 的最佳解决方案(使用我的代码 + 他的 map )快。相反，只使用普通数组。如果您注意到 S HashMap 的键是从 1 到 n' 或 n 的序列，您可以想出同样的想法div i 用于 i 从 1 到 r。所以我们可以将这样的 HashMap 表示为两个数组，根据搜索键在数组中进行查找。

我的代码 + Judy HashMap

$ time ./judy
95673602693282040

real    0m0.590s
user    0m0.588s
sys     0m0.000s

我的代码+我的稀疏图

$ time ./sparse
95673602693282040

real    0m0.203s
user    0m0.196s
sys     0m0.004s

如果不是 IOUArray 而不是已经生成的向量和使用 Vector 库并且 readArray 被 替换，这可以更快地完成不安全读取。但我认为，如果您对尽可能多地优化它并不真正感兴趣，则不应该这样做。

与此解决方案进行比较是作弊，是不公平的。我希望在 Python 和 C++ 中实现相同的想法会更快。但是@Krom 封闭 HashMap 的解决方案已经在作弊，因为它使用自定义数据结构而不是标准数据结构。至少您可以看到 Haskell 中标准和最流行的 HashMap 并没有那么快。使用更好的算法和更好的临时数据结构可以更好地解决此类问题。

Here's resulting code.