はじめに
Twitter で @mod_poppo さんが Haskellでもエラトステネスの篩がしたい！ というツ イートをしていました。
そこで こちらのコード に出てくる Data.Vector.Unboxed.ifoldr や Data.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 と実際の数は idxToNum と numToIdx で変換していきます。
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 が見つかります。
乗法より加法の方が処理が高速です。