Java の BitSet を使った Clojure の篩

昨日、Eclipse 環境で Clojure を使う準備をしたので、今日は簡単なプログラムを書いてみた。

簡単といえば Project Euler の1〜10番辺りかと見当をつけて眺めてみると、3番で素数を使っているのがあったので、エラトステネスのふるいでも書いてみることにした。

いろんな書き方があると思うけど、数ある関数型言語の中で敢えて Clojure を使うのは、やはり既存の Java 資産が使えるからというのが大きいと思うので、ここでは java.util.BitSet を使って実装してみることにした。

(ns euler.pe003)
(import 'java.util.BitSet)

(defn initial-bits [upper-bound] 
  (let [bits (BitSet. upper-bound)]
    (doall (map #(.set bits % true) (range 2 (inc upper-bound)))) bits))

(defn next-prime [pos upper-bound bits]
  (let [pos+ (inc pos)]
    (if (.get bits pos+) pos+ (next-prime pos+ upper-bound bits))))

(defn update-bits [pos upper-bound bits]
  (doall
    (map #(.set bits % false) (range (+ pos pos) (inc upper-bound) pos))) 
  [(next-prime pos upper-bound bits) bits])

(defn seive [n upper-bound bits sq] 
  (let [[m bits2] (update-bits n upper-bound bits)]
    (if (< sq m) bits2 (seive m upper-bound bits2 sq)))) 

(defn primes [upper-bound]
  (let [sq (int (Math/sqrt upper-bound))
        bits (seive 2 upper-bound (initial-bits upper-bound) sq)] 
    (filter #(.get bits %) (range upper-bound))))
 

実行してみると、以下のような結果が得られて、Haskell で計算してみたものと一致する。ただ 1,000,000 までの素数を数えているところでは、予想外に時間がかかっていた。もうちょい Clojure に慣れたら、なんで遅いのか調べてみたい。

=> (primes 100)
(2 3 5 7 11 13 17 19 23 29 31 37 41 43 47 53 59 61 67 71 73 79 83 89 97)
=> (count (primes 1000000))
78498

ここまでできたら、Problem 3 を解くのは簡単で、以下のように書き足して実行したら、正答が得られる。

(defn divmod [m n] [(long (/ m n)) (mod m n)])

(defn solve [a bits] 
  (let [[h & t] bits
        [q r] (divmod a h)]
    (if (not= r 0)
      (solve a t)
      (if (= 1 q) h (solve q bits)))))

(defn -main [& args]
  (solve 600851475143 (primes 10000)))
 

新しい言語で初めて書くプログラムとしては、当初の想定よりちょっと難しかったが、テキストを読んでるだけでは分からない事がいろいろあって、実際に書いてみるとやっぱり面白い。