Enumerable#sample_by_rate というメソッドでサンプリングするのはどうか

2014-10-14 02:00 追記

GitHub および RubyGems.org に sample_by_rate として公開しました。
https://github.com/yuya-takeyama/sample_by_rate
https://rubygems.org/gems/sample_by_rate

---

Ruby でデータのサンプリングといえば Array#sample というメソッドがあります。

(1..10).to_a.sample 3
# => [6, 1, 7]

もちろん結果は毎回異なります。

Enumerable ではなく Array のメソッドなので to_a が必要です。

この場合は to_a で解決することができましたが、以下のような場合は問題になります。

• データがとても大きい (to_a 時に多くのメモリを必要とする)
• 要素数がわからない (Array#sample の実装は、要素数がわかることを前提としているようです)
• 要素数が無限 (そもそも to_a ができない)

これは、Array#sample が「指定要素数分ランダムに取得する」という実装になっていることによる問題です。
これに対して、例えば「指定した確率で要素を yield する」というものであれば、要素数のわからない Enumerable や、Enumerator::Lazy による無限リストに対しても問題なく使えるでしょう。

というわけで作ってみた

module Enumerable
  def sample_by_rate(rate)
    return enum_for(__method__, rate) unless block_given?

    each {|n| yield n if rand < rate }
  end
end

使い方

以下は 1 から 1000 万までの数列を、10% にサンプリングして出力する例です。

(1..10000000).sample_by_rate(0.1) do |n|
  puts n
end

実行して行数をカウントすると、だいたい 10% ぐらいの数になっているのがわかります。

$ ruby sample.rb | wc -l
1001584

なかなか便利だと思うので gem にして公開するつもりです。
より良いメソッド名とかあればコメントいただけると幸いです。