[Ruby] ArrayじゃなくてEnumerableでsample（ランダムにn要素を選ぶ）を使えるようにする

enum_random.rb

#!/usr/bin/env ruby
# -*- coding: utf-8 -*-

# ArrayじゃなくてEnumerableでsample（ランダムにn要素を選ぶ）を使えるようにする。
# 例えば、大きいファイルから1行をランダムに抜き出すなど。
# 
# 厄介なのは、要素の個数が最初からはわからないこと。
# 作戦としては以下の通り。
# - k個目の要素までを読み込んだとき、そのk個目の要素が選ばれているべき確率は
#   n/k。なのでその確率で選ばれるかを決める。
# - 一度n個の要素に入っても、後から来たもので上書きされることはある。
#   上書きされる確率は等確率にする。

module Enumerable
  def sample(n = 1)
    n = n.to_i
    if n < 0
      raise ArgumentError, "Number of choices must be non-negative"
    elsif n == 0
      return []
    end
    
    k = 0
    stock = []
    each do |elem|
      k += 1
      if stock.size < n
        # まだn個選べていない場合は、単に選ぶ
        stock << elem
      else
        newpos = rand(k)
        stock[newpos] = elem if newpos < n
      end
    end
    
    stock
  end
end

if $0 == __FILE__
  if ARGV.size < 1
    STDERR.puts "[Usage of this file]"
    STDERR.puts "This library implements Enumerable#sample (almost the same as Array#sample, but random number generator cannot be specified)."
    STDERR.puts ""
    STDERR.puts "[Usage of this command-line tool]"
    STDERR.puts "This program samples [Count] lines in the file at random. If no file is specified, do the same for STDIN."
    STDERR.puts "Usage: enum_random [Count] ([FILES...])"
    exit
  end
  
  count = Integer(ARGV.shift)
  ARGF.each_line.sample(count)
end