Haskell
素数
エラトステネスの篩

Haskell で「エラトステネスの篩」 その1


はじめに

Twitter で @mod_poppo さんが Haskellでもエラトステネスの篩がしたい! というツ イートをしていました。

そこで こちらのコード に出てくる Data.Vector.Unboxed.ifoldrData.Vector.Unboxed.create といった関数を使って、数年前に作ったエラトステネ スの篩 ( 2 の倍数をあらかじめ除去したバージョン) を改良してみました。

さらに @mod_poppo さんに倣って、"2 または 3 の倍数をあらかじめ除去したバージョンのエラトステネスの篩" も作ってみました。

備忘録代わりにコードの解説もしてみたいと思います。


コード

import Control.Monad

import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

--
-- エラトステネスの篩 ( 2 の倍数をあらかじめ除去 )
--
primeList :: Int -> [Int]
primeList n = 2 : U.ifoldr f [] sieve
where
f i v xs = if v then idxToNum i : xs else xs

idxToNum i = 2 * i + 3
numToIdx v = div (v - 3) 2
lastIdx = numToIdx n

sieve = U.create $ do
mVec <- UM.replicate (lastIdx + 1) True
forM_ [0 .. numToIdx (floor $ sqrt $ fromIntegral n)] $ \i -> do
v <- UM.read mVec i
when v $ do
let num = idxToNum i
s = numToIdx (num * num)
forM_ [s, s + num .. lastIdx] $ \j -> do
UM.write mVec j False
return mVec

--
-- エラトステネスの篩 ( 2 または 3 の倍数をあらかじめ除去 )
--
primeList2 :: Int -> [Int]
primeList2 2 = [2]
primeList2 n = 2 : 3 : U.ifoldr f [] sieve
where
f i v xs = if v then num : xs else xs
where num = idxToNum i

idxToNum i = if (even i) then 6 * (div i 2) + 5 else 6 * (div i 2) + 7
numToIdx n = if (mod n 6 == 5) then (div n 6) * 2 else (div n 6) * 2 - 1
lastIdx
| mod6 < 1 = div n 6 * 2 - 2
| mod6 < 5 = div n 6 * 2 - 1
| otherwise = div n 6 * 2
where mod6 = mod n 6

sieve = U.create $ do
mVec <- UM.replicate (lastIdx + 1) True
forM_ [0 .. numToIdx (floor $ sqrt $ fromIntegral n)] $ \i -> do
v <- UM.read mVec i
when v $ do
let num = idxToNum i
s1 = numToIdx (num * num)
s2 = numToIdx (num * (idxToNum (i + 1)))
forM_ [s1, s1 + num * 2 .. lastIdx] $ \j -> do
UM.write mVec j False
forM_ [s2, s2 + num * 2 .. lastIdx] $ \j -> do
UM.write mVec j False
return mVec


"2 の倍数をあらかじめ除去したバージョン" の説明

primeList の中心となる sieve について書いていきます。

mVec <- UM.replicate (lastIdx + 1) True

mVec は「3 以上 n 以下の奇数 (3, 5, 7 ...)」を想定しています。 index と実際の数は idxToNumnumToIdx で変換していきます。

mVec の中の値は、素数ならば True 、合成数ならば False になります。

forM_ [0 .. numToIdx (floor $ sqrt $ fromIntegral n)] $ \i -> do

3 から「n の平方根以下」の奇数を順番に調べていきます。なぜ「n の平方根」まで調べれば十分かという説明は後程。

v <- UM.read mVec i

when v $ do
let num = idxToNum i
s = numToIdx (num * num)
forM_ [s, s + num .. lastIdx] $ \j -> do
UM.write mVec j False

v の値が True ならば num = idxToNum i は素数なので、 num の倍数を消していくのですが、 num の 2 乗から始めます。

なぜ num の 2 乗未満を無視できるかというと…例えば 7 の倍数を考えると、" $7 \times 3$ " は 3 の倍数として、" $7 \times 5$ " は 5 の倍数として既に消されているので、" $7 \times 7$ " 以降の数を調べればいいわけです。

前述の "「n の平方根」まで調べれば十分" という理由もここにあります。

num が「n の平方根」より大きい場合、 num の 2 乗は n より大きくなるので、これ以降の数は調べなくていいわけです。

7 の倍数は 7 跳びに出現するので、 numToIdx (7 * 7) , numToIdx (7 * 9) ... と計算しながら False を書き込む index を探していく必要はありません。

開始する index を s = numToIdx (7 * 7) で見つけたら、 s, s + 7, s + 7 + 7, s + 7 + 7 + 7 ... と 7 を足していくだけで False を書き込む index が見つかります。

乗法より加法の方が処理が高速です。


つづく…

Haskell で「エラトステネスの篩」その2