任意の階層だけ繰り返しをネストする多重 map 的メソッド(Ruby)

Array#nest_loop で、Integer#times の多重ネスト版です。配列にループ回数を入れて呼び出します。わかり切った多重ループを書くのが面倒なときに役立ちます。ブロックが与えられなければ Enumertor を返します。
 

[4, 2, 3].nest_loop do |i, j, k|
  print "#{[i, j, k]} "
end
#=>
[0, 0, 0] [0, 0, 1] [0, 0, 2]
[0, 1, 0] [0, 1, 1] [0, 1, 2]
[1, 0, 0] [1, 0, 1] [1, 0, 2]
[1, 1, 0] [1, 1, 1] [1, 1, 2]
[2, 0, 0] [2, 0, 1] [2, 0, 2]
[2, 1, 0] [2, 1, 1] [2, 1, 2]
[3, 0, 0] [3, 0, 1] [3, 0, 2]
[3, 1, 0] [3, 1, 1] [3, 1, 2]

 
for in 文を使って実装しているので、self の配列の中には範囲演算子や配列、ハッシュなどを使うことができます。また、ブロックのそれぞれの返り値が配列に入って返ります。なので、多重ループ版の map のように使えます。

p [2..6, 2, ["a", "b"]].nest_loop(&:itself)
#=>[[2, 0, "a"], [2, 0, "b"], [2, 1, "a"], [2, 1, "b"], [3, 0, "a"], [3, 0, "b"],
#   [3, 1, "a"], [3, 1, "b"], [4, 0, "a"], [4, 0, "b"], [4, 1, "a"], [4, 1, "b"],
#   [5, 0, "a"], [5, 0, "b"], [5, 1, "a"], [5, 1, "b"], [6, 0, "a"], [6, 0, "b"],
#   [6, 1, "a"], [6, 1, "b"]]

a = [:a, :b, :c]
p [a, a].nest_loop {|i, j| (j == :b) ? [i, "b"] : [i, j]}
#=>[[:a, :a], [:a, "b"], [:a, :c], [:b, :a], [:b, "b"],
#   [:b, :c], [:c, :a], [:c, "b"], [:c, :c]]



本体のコードです。(※注:2018/2/7 に修正しました。)
nest_loop.rb

class Array
  def nest_loop
    check = lambda do |obj|
      return 0...obj if obj.class.ancestors.include?(Integer)
      obj
    end
    
    e = Enumerator.new do |y|
      ns = lambda do |ax, args|
        a, ax = ax.first, ax.drop(1)
        for i in check.call(a)
          nxt = args + [i]
          if ax.empty?
            y << nxt
          else
            ns.call(ax, nxt)
          end
        end
      end
      ns.call(self, [])
    end
    block_given? ? e.each.with_object([]) {|i, h| h << yield(i)} : e
  end
end

 

Gem 化

自家製の utility Gem 'kaki-utils' に同梱しました。
kaki-utils | RubyGems.org | your community gem host
インストールは $ gem install kaki-utils で。使い方は

require 'kaki/utils/nest_loop'

a = [1, 2, 3]
p [a, a].nest_loop {|i, j| i * j}    #=>[1, 2, 3, 2, 4, 6, 3, 6, 9]

という感じ。(2018/2/6)