【Ruby のまずいコード】最初の n 個の素数

お題

最初の $n$ 個の素数を配列で返すメソッドを書いてください。
たとえば 4 を与えると [2, 3, 5, 7] を返します。
0 を与えると空配列 [] を返します。

素数を得るのは標準添付ライブラリー prime を使いましょう。

コード

コード 1

require "prime"

def prime_numbers(n)
  result = []
  count = 0
  Prime.each do |prime_number|
    break if count == n
    count += 1
    result << prime_number
  end
  result
end

コード 2

require "prime"

def prime_numbers(n)
  result = []
  Prime.each do |prime_number|
    return result if result.count == n
    result << prime_number
  end
end

講評

コード 1

コード 1 を再掲します。

require "prime"

def prime_numbers(n)
  result = []
  count = 0
  Prime.each do |prime_number|
    break if count == n
    count += 1
    result << prime_number
  end
  result
end

とても素朴なコードです。
いわゆるカウンター変数を使って，得られた素数の数を数えています。

こういうコードを書くときに注意すべきは，

• ループ脱出の条件分岐（break if count == n）
• カウントアップ（count += 1）
• ループ内のお仕事の本体（result << prime_number）

を書く順序です。これを間違うと正しい結果が得られません。
おそらく誰でも，n に具体的な小さな数（0，1，2）を当てはめて頭で流れを辿ってアルゴリズムを確認するのではないでしょうか。逆に言えばそういうシミュレーションをやらないと流れが把握しづらいコードなのです。

順序を変えて

require "prime"

def prime_numbers(n)
  result = []
  count = 0
  Prime.each do |prime_number|
    result << prime_number
    count += 1
    break if count == n
  end
  result
end

としてもおおむねうまくいきますが，n が 0 のときだけ無限ループに陥ります。
（脱出条件の count == n を count >= n とすれば無限ループは回避できますが，0 の場合の返り値は正しくありません）

おおむねうまくいくので，ミスに気づきにくいとも言えます。

p prime_numbers(4) # => [2, 3, 5, 7]

という結果を見て「よしよし，うまくいったぞ」と勘違いしがちです。

今回はお題の中に n が 0 の場合について言及しておいたので，「えっと 0 の場合は大丈夫かな？」と気にかけるでしょうが，ふだんのプログラミングでは引数の値の範囲をあまり意識しないでコードを書くことも多いので，気づきにくいバグが生まれます¹。

コード 2

コード 1 では，カウンター変数を使っていました。Ruby ではカウンター変数の出番はあまり多くありません。
カウンター変数を使っていたら「もっと簡潔にできないのかな？」と考えてみるといいと思います。

コード 2 はカウンター変数を使っていません。コード 2 を再掲します。

require "prime"

def prime_numbers(n)
  result = []
  Prime.each do |prime_number|
    return result if result.count == n
    result << prime_number
  end
end

要するに，「今できている配列が目的の長さになっていれば，その配列を返り値としてメソッドを脱出する」ということですね。
少しシンプルになったぶん，把握しやすくなりました。
しかし，こんなに何行も書かなければいけないものでしょうか？　そんなことはありません。

ちなみに，ループの中からいきなりメソッドを抜けていますが，大丈夫です。プログラミング言語によってはこれはまずいのですが，Ruby では何の問題もありません²。

改善

require "prime"

def prime_numbers(n)
  Prime.take(n)
end

と書けます。もはやメソッドにする必要すら感じないほど簡単になりました。

少しだけ解説します。

Prime はクラスであり，そのインスタンスが each メソッドを持っていて素数を順に取り出すことができるのですが³，いちいちインスタンスを参照しなくてもいいように，Prime クラスにクラスメソッド each が定義されていて，このクラスから直に素数を一つずつ取り出せるようになっています。

しかも Prime クラス自身が Enumerable なオブジェクトなので，Prime.each に基づいて Enumerable のメソッドが全部使えます⁴。

上のコードで，Prime.take(n) の take は Enumerable#take なのです。このメソッドは，与えられた回数だけ each を使って要素を取り出し，それを並べた配列を作って返すものです。

余談

お題がもし「$n$ までの素数を返す」のように上限を与えるものであれば，Prime.each に引数を与える用法で，初心者でも簡単に

require "prime"

def prime_numbers_up_to(n)
  result = []
  Prime.each(n) do |prime_number|
    result << prime_number
  end
  result
end

と書けますし，もう少し知識があれば

require "prime"

def prime_numbers_up_to(n)
  Prime.each(n).to_a
end

と非常に簡潔に書けますね⁵。

1. 逆に言えば，「ふだんからメソッドを書くときは想定すべき引数の値の範囲を意識しよう」となりますね。 ↩

2. ただし，こういうスタイルを嫌う向きもあります。とくにコードが長く複雑な場合，「中のほう」で return すると分かりづらいというのです。 ↩

3. Prime は（NilCLass などと同様）複数のインスタンスを作ることができず，Prime.new も使えないようになっています。Prime の唯一のインスタンスを Prime.instance で得ることができます。 ↩

4. とはいえ，Enumerable のメソッドの中には each が有限回で終わる前提のものも多く，こういったメソッドを Prime に対して使うと止まりません。二千何百年か前から知られているとおり，素数は無限にありますからね。いやもちろん，メモリー不足かなんかで止まります（落ちます）けど。 ↩

5. 後者のコードは，each のブロック無し用法と Enumerator の理解が必要です。 ↩