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

はじめに
========

Twitter で @mod_poppo さんが [Haskellでもエラトステネスの篩がしたい！](https://twitter.com/mod_poppo/status/1119915738364436482) というツ イートをしていました。

そこで [こちらのコード](https://gist.github.com/minoki/333fb9c2d82e2279f5bfda54b51b9279) に出てくる `Data.Vector.Unboxed.ifoldr` や `Data.Vector.Unboxed.create` といった関数を使って、数年前に作ったエラトステネ スの篩 ( 2 の倍数をあらかじめ除去したバージョン) を改良してみました。

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

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

コード
======

``` {.haskell}
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` について書いていきます。

``` {.haskell}
mVec <- UM.replicate (lastIdx + 1) True
```

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

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

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

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

``` {.haskell}
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 で「エラトステネスの篩」その２](https://qiita.com/little_Haskeller/items/a3ec4c82c7099f1c9cf1)