Alex Martelli版エラトステネスの篩をClojureで書いてみた

素数算出アルゴリズムで最も有名なものは多分エラトステネスの篩だろう。発見した素数の倍数を消していく（篩にかけていく）ことによって、残った数の中で最小のものが素数だとわかる。そのプロセスを繰り返していくことで、一定のn以下のすべての素数が効率よく算出できる。

Python界隈で素数について調べると、Alex Martelliによるエラトステネスの篩の無限版に行き着く。

Cooking with Python, Part 2 - O'Reilly Media

Alex Martelli自身もPython Cookbookの作者であり、StackOverflowやPython Mailing Listで非常に活発に答えていたので有名な人。さらにこのアルゴリズムは、import thisやtimsortで有名なTim Peters、itertoolsやdequeなどで有名なRaymond Hettingerなどが改良しており、Python Allstars的な感がある。

これをClojureに書き直してみたい。

とりあえず書いてみると、以下のような見るも無残なコードになった。

(defn next-composite [x p D]
  (if (contains? D x)
    (recur (+ x p) p D)
    x))

(defn sieve-step [[p D n]]
  (if (contains? D n)
    (let [p2 (D n)
          x (+ n p2)
          D2 (dissoc D n)
          x2 (next-composite x p2 D2)
          new-D    (assoc D2 x2 p2)]
      (recur [p new-D (inc n)]))
    [n (assoc D (* n n) n) (inc n)]))

(def primes (map first (iterate sieve-step [2 {4 2} 3])))

変数にわかりやすい名前をつけたいのだが、あまりいい案がなかった。そういう場合は、多分途中経過的な変数を作っているのがそもそもの間違いなんだと思う。

アルゴリズムの構造もお世辞にもわかりやすいとは言えないし。ループの中のループを別の関数にして、そこでも末尾再帰しているのは冗長ではないか。

というわけで以下のようにさらにリファクターしてみる。

(defn sieve-step [[x composites-map]]
  (let [n              (inc x)
        prime-factor   (composites-map n)]
    (if (nil? prime-factor)
      [n (assoc composites-map (* n n) n)]      
      (let [next-composite 
              (->> (iterate (partial + prime-factor) n)
                   (filter (complement composites-map))
                   first)
            new-comp-map
              (->  composites-map 
                   (dissoc n) 
                   (assoc next-composite prime-factor))]
        (recur [n new-comp-map])))))

(def primes (map first (iterate sieve-step [2 {4 2}])))

変数をごっそりカット（引数も無駄があったので整理）したことで名前が付けやすくなった。関数内ループをiterate/filter/firstで表現すれば末尾再帰な別関数を用意しなくてもいい。

これだと比較的流れがわかりやすい気がする。今まで出てきた素数からできる最小の合成数のmapを使って、もしある数がそのmapに含まれていなければ素数、含まれている場合は合成数。今まで出てきた素数一つにつきmapに必ず一つ対応する合成数が存在しているのがポイント。

Pythonだと効率はnaiveな書き方に比べて歴然としていたのだが、Clojureではどうだろう。

比較対象として、以下のコードを書いてみた。

(def odds-gte-3 (iterate #(+ 2 %) 3))

(defn prime? [n]
  (->> odds-gte-3
       (take-while #(< % (Math/sqrt (inc n))))
       (every? #(pos? (rem n %)))))

(def primes (concat [2] (filter prime? odds-gte-3)))

順次割って余が出ないことを確認する方式の素数ジェネレータ。奇数だけを平方根までの奇数で割って調べる、といった簡単な最適化を適用している。

これで

(time (apply + (take n primes)))

を比較してみたところ、

エラトステネスの篩：n = 10000 0.2秒、n = 100000 3秒、 n = 1000000 45秒

単純な割り算：n = 10000 0.3秒、n = 100000 12秒、 n = 1000000 7分前後

というわけで明確にスケールしている。