IndexErrorを起こさないArray#transpose

Ruby基本的なメソッドの話題です。

Rubyで配列の配列を行列と見立てて、転置行列を求める場合、Array#transposeを使いますが、このメソッドは、各要素数が異なるとIndexErrorが発生します。これを、要素数が異なっていても、後ろの要素にnilが入ってるとみなして、転置行列が作れるようにしたいというものです。

例を挙げた方が早いので、簡単な例です。入力となる配列の配列は下のようなもの。

data = [[1, 2, 3], [4], [5, 6, 7, 8]]

この配列の配列から、下の配列の配列を生成します。

[[1, 4, 5], [2, nil, 6], [3, nil, 7], [nil, nil, 8]]

最初のやり方: Array#transposeを使う

一つ目のやり方は、全ての要素の要素数を同じにしてから、Array#transposeを呼ぶものです。

max = data.max_by{|a| a.length}.length
data.each {|a| a[max-1] = nil if a.length < max}
p data.transpose # [[1, 4, 5], [2, nil, 6], [3, nil, 7], [nil, nil, 8]]

全ての要素の要素数を同じにする処理が、冗長な感じです。

2つ目のやり方: Enumerable#zipを使う

2つ目のやり方は、配列を組み合わせることができるEnumerable#zipを使って、生成するやり方です。このメソッドなら、Array#transposeメソッドと異なり、足りない要素はnilで埋められ、余分な要素は捨てられます。

ただし、レシーバの配列の要素数に揃えられます(↓)。

data[0].zip(*data[1..-1]) # [[1, 4, 5], [2, nil, 6], [3, nil, 7]] → NG([nil, nil, 8]も欲しい)

そこで、先頭の配列の要素数が、最大の要素数になるようにします。

max = data.max_by{|a| a.length}.length
data[0][max-1] = nil if data[0].length < max
p data[0].zip(*data[1..-1]) # [[1, 4, 5], [2, nil, 6], [3, nil, 7], [nil, nil, 8]]

こっちのがスッキリしていると思いますが、まだもっとシンプルに書ける気が…。うーん…。