Rubyで上位n件を求める

最近は仕事先以外でコードを書くことが多くなっている。
特にここ１・２ヶ月はパズルとかの問題を解くために実験的に多くの組み合わせの計算をRubyにさせることが多い。このときにしばしば欲しくなるのが「上位n件を求める」機能。配列に設定済みのデータであれば、Array#sortを呼び出した後に配列の先頭を見れば良いのだけど、以下の点が気になる。

• 上位10件を求めるときに100万件のデータを貯めておいて完全にソートするのはメモリと時間の無駄
• Arrayで扱わない無限個のデータには対応できない

例えば、一万個の乱数を求めてその値の上位5件を表示させたいのなら、こんな感じで簡単に取り出したい。（以下、コードはすべてRuby 1.9）

topn = TopN.new(5)
1.upto(10000).each do
  topn << Random.rand
end
topn.each_with_index do |data,i|
  puts "Rank#{i+1}:value=#{data}"
end

ついでに、個々のデータに付加情報があるのであれば、データから値を取り出す処理をコンストラクタにブロックとして与えればうまく動くようにしたい。例えば、「値が何回目に出現したか」を扱いたい時はこんな感じで書けたらなぁ...と。

topn = TopN.new(5) { |data| data[:value] }
1.upto(10000).each do |count|
  topn << {value: Random.rand,description: count}
end
topn.each_with_index do |data,i|
  puts "Rank#{i+1}:value=#{data[:value]} count=#{data[:description]}"
end

探してみても相当するライブラリが見つからなかったので、TopNクラスのコードを書いてみた。
（ちゃんと探せば実は有る？？）

最初の例を実行すると以下のように出力され、

Rank1:value=0.9999727375403678
Rank2:value=0.9999709710675848
Rank3:value=0.9998929520002925
Rank4:value=0.9998796385129795
Rank5:value=0.9998750822342006

二つ目の例であれば、以下のように出力される。

Rank1:value=0.9999711913114 count=3520
Rank2:value=0.9999658723206534 count=4642
Rank3:value=0.9999289720505727 count=9733
Rank4:value=0.9998648527153264 count=916
Rank5:value=0.9998544367073094 count=3624

数値に限らず、<や>で比較可能なオブジェクトであれば何でも対象に出来る。
以下、コード。

# 上位n件のソートされたデータ
class TopN
  attr_reader :top_data

  def initialize(n,&b)
    @n,@sel_v,@top_data = n,b,[]
  end

  def each
    @top_data.each { |v| yield v }
  end

  def each_with_index
    @top_data.each_with_index { |v,i| yield v,i }
  end

  def <<(target)
    update(target) if (@top_data.size < @n || to_v(@top_data[-1]) < to_v(target))
  end

  protected
  def to_v(v)
    @sel_v ? @sel_v.call(v) : v
  end

  def update(target)
    if @top_data.empty?
      @top_data << target
    else
      dest_data = []
      @top_data.each do |source|
        if target && to_v(target) > to_v(source)
          dest_data << target 
          target = nil
        end
        break if dest_data.size == @n
        dest_data << source
      end
      @top_data = dest_data
    end
  end
end