素数 - 倭算数理研究所

Scala で『Java によるアルゴリズム事典』のコードを書き換えてみようシリーズ（目次）。　ちょっと別用で素数 (prime number) を生成するコードを書いたので、こちらのブログで整理しておきます。　別用記事で見たように Int や Long に対してなら2行で書けますが、ここでは spire の Integral を使ってジェネリックなコードにしましょう。

アルゴリズム

ある自然数が与えられたとき、それが素数かどうかを判定するには、その自然数より小さい全ての素数で割りきれないかどうかを見ます。　実際には、その自然数の（正の）平方根以下の素数で割りきれなければ大丈夫です。　なぜなら、平方根以上の素数で割ったときに割り切れるなら、その商は平方根以下で、その素因数は平方根以下の素数になるからです。

実装

準備

まずはちょっとユーティリティメソッドを定義しておきましょう。　標準 API の Stream.from() メソッドを使えば、無限に続く等差数列を与える Stream オブジェクトを簡単に作れますが、spire の Integral 型に対して同様のメソッドがないので自作しておきます*1：

import spire.math.Integral
import spire.implicits._

object IntegralSequence{
  def from[I: Integral](start: I): Stream[I] = from(start, 1)
  def from[I: Integral](start: I, step: I): Stream[I] = start #:: from(start + step)
}

#:: は Stream トレイトに定義されているメソッドで（「:」で終わるので右結合）、遅延評価で無限 Stream を生成します。　Stream を再帰的に定義するときなどに使います。　詳しくは Stream の Scaladoc を参照のこと。

使い方は Stream.from() メソッドと同じです。　以下では引数を2つとる方しか使いませんが、なんとなく引数1つのもの（公差1の等差数列を生成）も定義しています。

素数を生成するコード

では、本題の素数を生成するコード。　以下では4つのメソッドを定義していますが、素数の Stream を生成するアルゴリズムに関係するのは後の2つ。　ただ、簡単に使うには前の2つのメソッドの方が便利かと思います：

import spire.math.Integral
import spire.implicits._

object PrimeNumber{

  /** n番目の素数を返す（素数2は0番目） */
  def apply[I: Integral](n: I): I = streamOf.apply(n.toInt)

  /** Int 型の素数列を返す */
  def stream: Stream[Int] = streamOf[Int]

  /** I 型の素数列を返す */
  def streamOf[I: Integral]: Stream[I] =
    2 #:: IntegralSequence.from[I](3, 2).filter(isPrime(_))

  /** 素数かどうかを判定する */
  def isPrime[I: Integral](n: I): Boolean =
    n > 1 && streamOf.takeWhile(p => p*p <= n).forall(n % _ != 0)
}

apply() メソッドは引数で指定された番号の素数を返します。　素数2は0番目です。
stream メソッドは Int の Stream を返します。
streamOf メソッドは型パラメータに指定した型の Stream を返します。　たとえば Long 型の Stream を返す場合は streamOf[Long] とします。
- IntegralSequence.from() メソッドは「準備」のところで定義したメソッドで、ここでは3以上の奇数（3で始まり２ずつ増える）を返す Stream を生成します。　偶数は2以外に素数がないので試すだけ無駄ですね。
- #:: メソッドは準備のところでも出てきた無限 Stream を生成するメソッドで、ここでは一見再帰になってませんが、isPrime() が streamOf() を呼び出しているので遅延評価にする必要があります。
- このコードでは、最小で唯一偶数の素数2を先頭にして、その後に奇数から isPrime によって素数をフィルタリングして返しています。
isPrime() メソッドは引数の整数が素数かどうかを判定します。　streamOf で返される素数列のうち、引数の平方根より小さいものだけを取り出して、それらすべての素数が引数を割り切らないかどうかを判定しています。

streamOf も isPrime も実質1行ずつで書けてるので、2行で書けているといってもいいでしょう（棒）　ちなみに、Int 型に関して2行で書くとこんな感じです：

  val primes: Stream[Int] = 2 #:: Stream.from(3, 2).filter(isPrime)
  def isPrime(n: Int): Boolean = n > 1 && primes.takeWhile(p => p*p <= n).forall(n % _ != 0)

サンプル・コード

では、上記の（ジェネリックな方の）コードを動かしてみましょう。

// apply() メソッドで n 番目の素数を取得する
assert( PrimeNumber(0) == 2 )
assert( PrimeNumber(10) == 31 )

// 素数の Stream[Int] を返す
assert( PrimeNumber.stream.take(10) == Seq(2, 3, 5, 7, 11, 13, 17, 19, 23, 29) )

// 素数の Stream[Long] を返す
assert( PrimeNumber.streamOf[Long].take(10) == 
  Seq(2L, 3L, 5L, 7L, 11L, 13L, 17L, 19L, 23L, 29L) )