zdogma's diary

徒然なるエンジニアな日々。

1からNまでの素数の数を数える標準ライブラリ Prime を読んでみる

ことはじめ

とある機会に「1からNまでの素数の数を数えよ」という問題を解こうとした時に、最初に下記のように書いた。

(2..number).each_with_object([]) { |num, primes|
  sqrt_primes = primes.select { |prime| prime < Math.sqrt(number) }
  primes << num if sqrt_primes.none? { |prime| num % prime == 0 }
}.count

ベンチマークを取れば歴然だが、下記のように N が大きくなると処理時間が大きくなる。 結果、テストケースを所定の時間内にクリアできないという状態に。。。

                       user     system      total        real
number = 10             0.000000   0.000000   0.000000 (  0.000045)
number = 100            0.000000   0.000000   0.000000 (  0.000342)
number = 1000           0.010000   0.000000   0.010000 (  0.012887)
number = 10000          0.660000   0.000000   0.660000 (  0.671416)

ぱっとチューニングのアイデアが思い浮かばなかったので、素数判定のロジックを調べていると、どうやら標準ライブラリに Prime というものがあるらしい。 下記のようなコードを書いてベンチマークを取ってみると...

require 'prime'
numbers.each { |number| Prime.each(number).count }
                       user     system      total        real
number = 10             0.000000   0.000000   0.000000 (  0.003707)
number = 100            0.000000   0.000000   0.000000 (  0.003185)
number = 1000           0.010000   0.000000   0.010000 (  0.002970)
number = 10000          0.000000   0.000000   0.000000 (  0.002980)

こちらだと余裕でテストケースをクリア。 なんとも悔しい気持ちになったので、Prime の実装を読んでみることにした。

本題

Prime の実装は下記にある。

ruby/prime.rb at trunk · ruby/ruby · GitHub

Prime では、下記のように each の引数に与えた数字までの素数を返す。

[56] pry(main)> Prime.each(100).to_a
=> [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]

[57] pry(main)> Prime.each(100).class
=> Prime::EratosthenesGenerator

素数の判別方法にはいろいろなやり方があるが、Prime でデフォルトで採用されているのは「エラトステネスの篩(ふるい)」と呼ばれているアルゴリズムである。 wikipedia によると、下記のように実装されるらしい。

  • ステップ 1
    • 探索リストに2からxまでの整数を昇順で入れる。
  • ステップ 2
    • 探索リストの先頭の数を素数リストに移動し、その倍数を探索リストから篩い落とす。
  • ステップ 3
    • 上記の篩い落とし操作を探索リストの先頭値がxの平方根に達するまで行う。
  • ステップ 4
    • 探索リストに残った数を素数リストに移動して処理終了。

エラトステネスの篩 - Wikipedia

詳細な説明は省くが、実際に該当のアルゴリズムを実装しているのは下記の部分。

def compute_primes
  # max_segment_size must be an even number
  max_segment_size = 1e6.to_i
  max_cached_prime = @primes.last
  # do not double count primes if #compute_primes is interrupted
  # by Timeout.timeout
  @max_checked = max_cached_prime + 1 if max_cached_prime > @max_checked

  segment_min = @max_checked
  segment_max = [segment_min + max_segment_size, max_cached_prime * 2].min
  root = Integer(Math.sqrt(segment_max).floor)

  segment = ((segment_min + 1) .. segment_max).step(2).to_a

  (1..Float::INFINITY).each do |sieving|
    prime = @primes[sieving]
    break if prime > root
    composite_index = (-(segment_min + 1 + prime) / 2) % prime
    while composite_index < segment.size do
      segment[composite_index] = nil
      composite_index += prime
    end
  end

  @primes.concat(segment.compact!)

  @max_checked = segment_max
end

おわりに

最近競技プログラミングの勉強を始めているが、やはり数学への理解がかなり重要で、直感で消費メモリ・処理時間がイメージできるまで頭を慣らしたい。
今回はとりあえずアルゴリズムの内容がわかったので、次回、自分のアルゴリズムとの差分と処理時間の差が生まれた要因をもう少し突き詰めてみようと思う。