Rubyの配列の要素をn個の配列へ順繰りに振り分け

配列や範囲オブジェクトなどに対し以下のように動作する Enumerable#balance を実装したい。ロードバランサーのラウンドロビン方式みたいな感じ。

pp (:A..:Z).balance(7) # 均等にならない場合
#=> [[:A, :H, :O, :V],
#    [:B, :I, :P, :W],
#    [:C, :J, :Q, :X],
#    [:D, :K, :R, :Y],
#    [:E, :L, :S, :Z],
#    [:F, :M, :T],
#    [:G, :N, :U]]

pp (:Α..:Ω).balance(4) # 均等になる場合
#=> [[:Α, :Ε, :Ι, :Ν, :Ρ, :Φ],
#    [:Β, :Ζ, :Κ, :Ξ, :Σ, :Χ],
#    [:Γ, :Η, :Λ, :Ο, :Τ, :Ψ],
#    [:Δ, :Θ, :Μ, :Π, :Υ, :Ω]]

pp [].balance(3) # 空配列の場合
#=> [[], [], []]

---

作ったらこうなった。

module Enumerable
  def balance(n)
    m = nil
    each_slice(n).to_a.tap do |*, la|
      return Array.new(n) { [] } unless la
      m = la.size
      la[n-1] = la[n-1] # nil padding
    end.transpose.tap do |ta|
      ta.drop(m).each(&:pop) # remove padding
    end
  end
end

• 基本戦略は each_slice(n).to_a.transpose 。
• しかし、転置 Array#transpose は各要素の大きさが揃っていることが必要。
• なので転置の前後でパディングの追加と削除を実施する。
  • 転置前は最終行の配列長を他と同じになるよう増やす。（参考：不足分をnilで埋めた配列を取得）
  • 転置後は増やした分だけ各行の末端を削る。（※ #compact だと元々あったnilまで消してしまう）
• 細かい配列操作は破壊的に行い、なるべく新しいインスタンス生成を抑える。
  • Object#tap が活躍。
• 空配列の転置は都合が悪いので、早めに検知してreturnする。（参考：空の配列に対するRubyのメソッドの挙動）

---

無限のリストでも処理できるようにn個のEnumeratorを返すことも考えたが、使い勝手や効率がよくわからないので保留。

module Enumerable
  def lazy_balance(n)
    enum = lazy
    Array.new(n) { |i| enum.drop(i).each_slice(n).map(&:first) }
  end
end

これは Enumerator::Lazy の配列を返しているため、最後に各要素を #force （ #to_a ）すれば最初の例と同じ結果が得られる。

pp (:A..:Z).lazy_balance(7).map(&:force)

pp (:Α..:Ω).lazy_balance(4).map(&:force)

pp [].lazy_balance(3).map(&:force)