Ruby Set を Hash で置き換える - インゲージ開発者ブログ

こんにちは、masm11 です！

みなさんは、Ruby で集合を表現する時に、何を使いますか？ Set を使いますか？ Set って内部的には Hash を使ってますよね。そう、Hash でも実現できます。

今回は、Set を Hash で置き換えた話をします。

Set とは

Set とは、ようするに集合です。集合なので、同じ要素は2つ以上ありません。

require 'set'

set = Set.new
(1..10).each do |i|
  set << i
end
(1..10).each do |i|
  set << i
end

set.each do |i|
  puts i
end

luna:~ % ruby set1.rb
1
2
3
4
5
6
7
8
9
10

集合なので、時々、交差を取りたくなります。交差、intersection。つまり、両方の集合に存在する要素を得る演算です。

set1 = Set.new
(1..10).each do |i|
  set1 << i
end

set2 = Set.new
(5..15).each do |i|
  set2 << i
end

set = set1.intersection(set2)
set.each do |i|
  puts i
end

これを実行すると、set1 には 1～10、set2 には 5～15 が入っているので、その intersection は 5～10 です。

luna:~ % ruby set2.rb
5
6
7
8
9
10

intersection の実装

ところで、intersection はどのように実装されているのでしょう？

実装はこちらにあります。

  def &(enum)
    n = self.class.new
    if enum.is_a?(Set)
      if enum.size > size
        each { |o| n.add(o) if enum.include?(o) }
      else
        enum.each { |o| n.add(o) if include?(o) }
      end
    else
      do_with_enum(enum) { |o| n.add(o) if include?(o) }
    end
    n
  end
  alias intersection &

n は結果ですね。今回は set1 も set2 も Set なので、enum.is_a?(Set) は true です。そして、サイズは同じなので、else です。結局、

        enum.each { |o| n.add(o) if include?(o) }

↑この一行がすべてのようです。

引数の要素すべてについて、それがレシーバにもあれば、それを n に追加する

なるほど。include? はどのように実装されているのでしょう？

  def include?(o)
    @hash[o]
  end

そのまんまですね。@hash は、キーは要素で、値は true です。

Ruby でループしてるんですね… これって遅くないですか？

Array の場合と比べてみる

とりあえず、遅そうな Array の場合と比べてみます。

N = 10000000

set1 = Set.new(1..N)
set2 = Set.new(1..N)

ary1 = (1..N).to_a
ary2 = (1..N).to_a

Benchmark.bm do |x|
  x.report {
    set1.intersection(set2)
  }
  x.report {
    ary1.intersection(ary2)
  }
end

結果がこちら。

luna:~ % ruby set_vs_array.rb
       user     system      total        real
   5.815057   0.184966   6.000023 (  6.000748)
   2.950291   0.135012   3.085303 (  3.087259)

Array の方が倍速！

Array の方は C で書かれているとは言え、速すぎませんか？要素数1千万の intersection ですよ？

一体どんな秘密が隠されてるのでしょうか？ C の実装を見てみました。

    hash = ary_make_hash(ary2);

    for (i=0; i<RARRAY_LEN(ary1); i++) {
        v = RARRAY_AREF(ary1, i);
        vv = (st_data_t)v;
        if (rb_hash_stlike_delete(hash, &vv, 0)) {
            rb_ary_push(ary3, v);
        }
    }

ary2 から hash を作る
ary1 各要素について、hash に要素があれば、それを削除して ary3 (結果) に突っ込む

なるほど… Hash を作ることで高速化してました。これなら速いですね。

`Set#intersection` を再実装する

Set は遅いです。Array にすら負けています。というか、Ruby でループしてる時点で負けでしょう。

Set をやめて Hash を直接使うことにします。 Hash なら C で実装されてるので速いはず！

hash1 = Hash[ary1.zip([true] * N)]
hash2 = Hash[ary2.zip([true] * N)]

ここまで書いて、はたと気づきました。Hash には intersection なんてないんですよ… どうしましょう… 何か代わりになるメソッドは…

これですね、slice。引数にキーを並べればいいようです。

N = 10000000

set1 = Set.new(1..N)
set2 = Set.new(1..N)

ary1 = (1..N).to_a
ary2 = (1..N).to_a

hash1 = Hash[ary1.zip([true] * N)]
hash2 = Hash[ary2.zip([true] * N)]

Benchmark.bm do |x|
  x.report {
    set1.intersection(set2)
  }
  x.report {
    ary1.intersection(ary2)
  }
  x.report {
    hash1.slice(*hash2.keys)
  }
end

実行結果が↓こちら。

luna:~ % ruby by-hash.rb     
       user     system      total        real
   6.361795   0.188878   6.550673 (  6.554033)
   3.133403   0.129874   3.263277 (  3.265233)
 by-hash.rb:25:in `block (2 levels) in <main>': stack level too deep (SystemStackError)
        from /usr/lib/ruby/3.0.0/benchmark.rb:296:in `measure'
        from /usr/lib/ruby/3.0.0/benchmark.rb:378:in `item'
        from by-hash.rb:24:in `block in <main>'
        from /usr/lib/ruby/3.0.0/benchmark.rb:176:in `benchmark'
        from /usr/lib/ruby/3.0.0/benchmark.rb:208:in `bm'
        from by-hash.rb:17:in `<main>'
luna:~ %

あらら… どうも、スタックが溢れたようです…。こういう時は RUBY_THREAD_VM_STACK_SIZE を指定すれば大きくすることができます。もう一度！

luna:~ % RUBY_THREAD_VM_STACK_SIZE=100000000 ruby by-hash.rb 
       user     system      total        real
   6.086419   0.187998   6.274417 (  6.275484)
   2.960836   0.119998   3.080834 (  3.081353)
   3.381075   0.120012   3.501087 (  3.501670)

Array には負けましたが、Set よりは十分高速ですね！

メモリは食いますが、今回、N=10000000 なので仕方ないですね。 N=10000 くらいなら環境変数なしでも全然問題ありません。