Life Goes On

まあまあだけど楽しんでる方です

速い素数の求め方

Haskell

ちょっと前に haskell-cafe で速い素数の求め方が話題になっていました。
どんなもんだか確認した結果を記録。
僕がよく使ってたコードはこんな感じ。改めて見るとツッコミどころ満載です。
106 番目の素数を求めるのに手元のマシンで 2 分以上かかります。

primes0 = 2 : filter isPrime0 [3..]
isPrime0 x = all ((/=0).mod x) $ takeWhile (<= (floor$sqrt$fromIntegral x)) primes0
  • 平方根または二乗を求めるのに、いちいち整数から浮動小数に変換するのはコスト高すぎ
  • 偶数も全てチェック対象にしているのは無駄
  • mod より rem の方が速い（これはなぜ？）

というあたりを修正したのが↓のコード。1 分 20 秒。

primes1 = 2 : 3 : filter isPrime1 [5,7..]
isPrime1 x = all ((/=0).rem x) $ takeWhile (\y -> y*y<=x) primes1

偶数だけでなく 3 の倍数もチェック対象から外したら、というのはよく話題になりますが、あまり速くなりません。たぶんそういう数列（内包表記の部分）を生成するのにコストがかかるのでしょう。1 分 20 秒。

primes2 = 2 : 3 : 5 : filter isPrime [ 6*i+j | i<-[1..], j<-[1,5] ] where
  isPrime x = all ((/=0).rem x) $ takeWhile (\y -> y*y<=x) $ drop 2 primes2

という訳で、アルゴリズムを考え直します。
32 未満の数が素数かどうかをチェックするには、2 で割り切れるかどうかを調べれば十分です。
同様に 52 未満の数が素数かどうかをチェックするには、2, 3 で割り切れるかどうかを調べれば十分です。（5 以上の素数で割ってみる必要はありません）
つまり n 番目の素数の二乗未満までは、最初の素数 n-1 個で割り切れるかを調べればよいわけです。
これを愚直に実装したのが以下のコード。特に何も効率化していませんが、50 秒くらいで終わります。
（以前のコードと同じようなフィルタリングですが、判定の回数が減っています）

primes3 = 2 : sieve 0 [3..] where
  sieve i xs = filter isPrime as ++ sieve (i+1) bs where
    (ps,p:_) = splitAt i primes3
    (as,bs) = span (<p*p) xs
    isPrime x = all ((/=0).rem x) ps

さらに偶数をチェック対象から除くと、40 秒くらいに縮みます。

primes4 = 2 : sieve 0 [3,5..] where
  sieve i xs = filter isPrime as ++ sieve (i+1) bs where
    (ps,p:_) = splitAt i primes4
    (as,bs) = span (<p*p) xs
    isPrime x = all ((/=0).rem x) ps

関数 sieve の引数にリストを渡すよりは、数を渡して内部でリストを生成するほうが、速くなります。
Haskell Wiki に素数についてのページがありますが、その IV 番目に相当します。
30 秒。今のところ、これが最終形です。
ちなみに 32bit で収まるなら、primes5 :: [Int] と明示的に宣言した方が速くて 20 秒で終わります。（コンパイル時に Unbox 化されるらしい）

primes5  = 2 : primes' where
  primes' = 3 : sieve 0 5
  sieve i x = filter isPrime [x,x+2..p*p-2] ++ sieve (i+1) (p*p+2) where
    (ps,p:_) = splitAt i primes'
    isPrime x = all ((/=0).rem x) ps

おまけ。Haskell らしさを追及して高階関数で書くと以下のようになります。
人によってはこの方が分かりやすいかもしれません。

primes6 :: [Int]
primes6  = 2 : primes' where
  primes' = 3 : (concat $ snd $ mapAccumL sieve 5 [0..])
  sieve x i = (p*p+2, filter isPrime [x,x+2..p*p-2]) where
    (ps,p:_) = splitAt i primes'
    isPrime x = all ((/=0).rem x) ps

もっと速い求め方があれば、ぜひぜひコメントを。