汎用的でけっこう速い素数無限列挙 - 名古屋で数学するプログラマ(仮)

二週間ぶりです。風邪引いてました。二週間かかってようやく落ち着きました。この時期の夏風邪ってば以下略。

さてさて、今日から暫く「素数列挙」の話をします。

Ruby1.9 の素数列挙メソッド

Ruby1.9には、素数列挙・素数判定・素因数分解を一手に扱う素数ライブラリが用意されています。
このブログでも素数列挙する時とかに何の説明もなく使ってました。

require 'prime'
Prime.each {|pr|puts pr}
# => 2
# => 3
# => 5
# => 7
# => 11
# => …（以降無限に列挙）

これ、割と速いです。というか結構速いです。
python+素数列挙でググるとよく出てくる、こんなサンプル

from itertools import ifilter, count
def prime_generator():
  g = count(2)
  while True:
    prime = g.next()
    yield prime
    g = ifilter(lambda x, prime=prime: x%prime, g)

for pr in prime_generator():
  print pr
# => 2
# => 3
# => 5
# => 7
# => 11
# => …（以降無限に列挙）

よりもよっぽど速いです*1。その理由もちゃんとあるのですが、ここでは触れる程度にしておきます。

結構速いのですが、色々調べてみると、もっと速くできるらしいことが分かります。
特に「指定した数までの素数を列挙する」ことに特化すれば、かなり速くできます。例えば
エラトステネスのふるいをコンパクトにする方法 - Smalltalkのtは小文字ですとか、[ruby-list:16645] (http://blade.nagaokaut.ac.jp/cgi-bin/scat.rb/ruby/ruby-list/16645)とかを参照のこと。
でも、今回は「無限列挙」。しかも、Ruby1.9のprimeライブラリよりも高速なもの。しかもできれば、Prime#eachによる列挙と互換性を持たせたい*2。

Ruby1.9 の素数列挙メソッドの仕組み

Rubyのライブラリリファレンスを見れば分かるのですが、RubyのprimeライブラリのPrime#eachメソッドには、第2引数もあります。
$RUBY_HOME/lib/ruby/1.9.1/prime.rb の中を覗いてみると、こんな感じになっています（かなり抜粋）。

class Prime
  : #中略
  def each(ubound = nil, generator = EratosthenesGenerator.new, &block)
    generator.upper_bound = ubound
    generator.each(&block)
  end
  : #中略
end

EratosthenesGeneratorというのが肝。Primeクラスの内部クラスになっています。やはりprime.rbから定義を見てみます。

class Prime
  : #中略
  class PseudoPrimeGenerator
    include Enumerable
    : #中略
    def succ
      raise NotImplementedError, "need to define `succ'"
    end
    : #中略
  end

  class EratosthenesGenerator < PseudoPrimeGenerator
    : #中略
    def succ
      @last_prime = @last_prime ? EratosthenesSieve.instance.next_to(@last_prime) : 2
    end
    : #中略
  end
  : #中略
end

つまりPseudoPrimeGeneratorというクラスがあって、所謂抽象クラス(Abstract Class)で、EratosthenesGeneratorはそれを継承したクラスで、succメソッドを実装してやっている（そして上の抜粋コードでは省略しましたが、PseudoPrimeGenerator#eachメソッドはPseudoPrimeGenerator#succメソッドを利用して実装されている）、という仕組みになっています。

なので。
オリジナルの素数ジェネレータをPrime::PseudoPrimeGeneratorクラスを継承して作って*3 *4、そのインスタンスをPrime#eachメソッドの第2引数に渡せば、独自実装による素数列挙であってしかもprimeライブラリのPrimeクラスを拡張した（＝同じAPI仕様で動く）ものができる、というわけですね。

五月雨

ということで、そのようなクラスを1つ作ってみました。
メモ: 上限を決めずに素数を生成の高速化(Ruby版)を参考に、それを改良してさらに高速かつコンパクトに。

使用例。

require 'prime'
require 'benchmark'
require_relative 'prime.a2'

max = ARGV.size == 1 ? ARGV[0].to_i : 100
pr = 0
Prime.each(max, Prime::A2Generator.new) {|num| pr = num}
puts "Max prime num(<=#{max}) is #{pr}."
block = Proc.new {}

Benchmark.bmbm do |b|
  b.report("Prime.each(max)") {Prime.each(max, &block)}
  b.report("Prime.each(max,a2)") {Prime.each(max,Prime::A2Generator.new,&block)}
end

↑の実行例↓。

$ ruby prime_a2_benchmark.rb 600000
Max prime num(<=600000) is 599999.
Rehearsal ------------------------------------------------------
Prime.each(max)      1.680000   0.010000   1.690000 (  1.688865)
Prime.each(max,a2)   0.340000   0.000000   0.340000 (  0.347410)
--------------------------------------------- total: 2.530000sec

                         user     system      total        real
Prime.each(max)      0.330000   0.000000   0.330000 (  0.331884)
Prime.each(max,a2)   0.340000   0.000000   0.340000 (  0.341556)

うん、かなり速いです。少なくとも1回目は約5倍速です。
あれ、でも2回目はほとんど同じ数字をたたき出しています。というか数字の上では微妙に標準のPrime#eachの方が速いですね。これはなぜでしょう？
実は標準で使用されるPrime::EratosthenesGeneratorは、ある仕組みで生成した素数をキャッシュしているのです。2回目以降に呼ばれたときにはそのキャッシュを利用するので1回目よりもかなり速くなるのです。その結果、今回作ったPrime::A2Generatorと同じくらい速くなってしまっているのです。

じゃ、A2Generatorもキャッシュを持たせれば速くなる？
それも一つのアプローチですが、ここではそっちには進まず、違うアプローチに進みたいと思います。
キーワードは、"EratosthenesSieve"。
ということで今回はここまで。続きは次回に。

*1:なぜRubyの話なのにpythonの例を出したのかというと、この例そのままをRubyで同じような動きをするモノを書くのが意外と面倒だったためです。言語レベルで「どっちが優れているか」という話ではないことを言及しておきます。

*2:最後のはもともと、著者によるワガママですけれど。

*3:先の抜粋コードで省略してしまっているのですが、このときsuccメソッドだけでなく、next、rewindの2つのメソッドも実装する必要があります。

*4:というかRubyなので、upper_bound=メソッドとeachメソッドの2つを実装したクラスであればPrime::PseudoPrimeGeneratorクラスを継承したクラスでなくても実はOKです。