RubyでN次元の超直方体中の格子点に対してループを回す

問題設定

２次元の長方形（xの領域は [x0,x1] , yの領域は [y0,y1]）の中に含まれる格子点に対してループを回したい場合、Rubyでは以下のように書くことができる。
（ただし、x0,x1,y0,y1はIntegerとする）

(x0..x1).each do |x|
  (y0..y1).each do |y|
    puts “#{x}, #{y}"
  end
end

３次元の直方体の場合は、同様に３重ループを使って

(x0..x1).each do |x|
  (y0..y1).each do |y|
    (z0..z1).each do |z|
    puts “#{x}, #{y}, #{z}"
  end
end

これを一般のN次元の場合にRubyでどう書けばよいか？

実装方法１ Array#productを使う

Arrayには直積を計算するproductというメソッドが存在する。これを用いると以下のように簡潔に書くことができる。

ruby

ranges = [ 1..3, 4..5, 6..9 ]  # 各次元の定義域をRangeで持つ配列

arrays = ranges.map {|range| range.to_a }
arrays[0].product( *arrays[1..-1] ) {|a| p a }   # => [1,4,6], [1,4,7], .... [3,5,9]

ただし、productにblockを渡さない場合は格子点をすべて含んだ配列を作成してしまう
Nが大きくなって格子点の数が指数関数的に増えてきた場合、メモリ消費量やオブジェクトの生成コストが膨れ上がってしまうので注意。
blockを渡していれば格子点に対してループを回しながらブロックを評価する。

実装方法２ 自前で再帰的に書く

再帰的にループにする。Enumeratorを使った。

ruby

ranges = [ 1..3, 4..5, 6..9 ]

e = Enumerator.new do |y|
  point = Array.new( ranges.size, 0 )

  iterate_in = lambda do |d|
    ranges[d].each do |x|   # d次元目に対するループ
      point[d] = x
      if d + 1 < point.size
        iterate_in.call( d+1 )
      else
        y << point
      end
    end
  end

  iterate_in.call(0)
end

e.each {|(x,y,z)| p [x,y,z] }    # => [1,4,6], [1,4,7], ..... [3,5,9]

この方法の場合、eはEnumeratorのオブジェクトなので each_with_indexなどのメソッドも利用出来る。

e.each_with_index {|(x,y,z),idx| p [x,y,z],idx }
# =>
#   [1, 4, 6]
#   0
#   [1, 4, 7]
#   1
#   ...

二つの方法の比較

速度の差が無いか確認してみた。

• 実行環境
  • ruby 2.2.0p0
  • Mac Pro (Mid2012) (3.2GHz Quad-Core Intel Xeon)

実行時間の測定にはBenchmarkモジュール、メモリ使用量の測定には ObjectSpace.memsize_of_all を利用した。

test_compare.rb

require 'objspace'
require 'benchmark'

def points(ranges)
  Enumerator.new do |y|
    point = Array.new( ranges.size, 0 )

    iterate_in = lambda do |d|
      ranges[d].each do |x|
        point[d] = x
        if d + 1 < point.size
          iterate_in.call( d+1 )
        else
          y << point
        end
      end
    end

    iterate_in.call(0)
  end
end

def points2(ranges)
  arrays = ranges.map {|range| range.to_a }
  arrays[0].product( *arrays[1..-1] ) {|a| yield a }
end

memusage = {before: ObjectSpace.memsize_of_all }
ranges = [1..10, 1..10, 1..10, 1..10, 1..10, 1..10, 1..10]  # 10^7点に対するiteration
Benchmark.bm 10 do |r|
  r.report "enumerator" do
    points(ranges).each {|(x,y,z)| x+y+z }
    memusage[:enumerator] = ObjectSpace.memsize_of_all
  end

  r.report "product" do
    points2(ranges) {|(x,y,z)| x+y+z }
    memusage[:product] = ObjectSpace.memsize_of_all
  end
end
memusage.each {|key,val| memusage[key] = "#{val/1048576}MB" }
p memusage

実行結果は

result2.txt

                 user     system      total        real
enumerator   3.550000   0.010000   3.560000 (  3.559787)
product      4.780000   0.030000   4.810000 (  4.834131)
{:before=>"2MB", :enumerator=>"2MB", :product=>"2MB"}

• productはCで実装されているので速いかと思ったが、少しだけenumeratorを使った方が速かった。Enumeratorの実装では、ループ変数を使いまわしているのでオブジェクトの生成コストが抑えられているのかもしれない。
• メモリ消費についてはどちらの実装でも適切に抑えられている。試しにproductに対してblockを渡さずに実行したところ1GB近くのメモリを消費していた。