implement Array.transpose

2 files changed, 39 insertions, 0 deletions

diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index b525d3006..c0995bb99 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -904,4 +904,32 @@ class Array
       end
     end
   end
+
+  ##
+  # call-seq:
+  #    ary.transpose -> new_ary
+  #
+  # Assumes that self is an array of arrays and transposes the rows and columns.
+  #
+  # If the length of the subarrays don’t match, an IndexError is raised.
+  #
+  # Examples:
+  #
+  #    a = [[1,2], [3,4], [5,6]]
+  #    a.transpose   #=> [[1, 3, 5], [2, 4, 6]]
+
+  def transpose
+    return [] if empty?
+
+    column_count = nil
+    self.each do |row|
+      raise TypeError unless row.is_a?(Array)
+      column_count ||= row.count
+      raise IndexError, 'element size differs' unless column_count == row.count
+    end
+
+    Array.new(column_count) do |column_index|
+      self.map { |row| row[column_index] }
+    end
+  end
 end
diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb
index bf96b885e..7d810acc2 100644
--- a/mrbgems/mruby-array-ext/test/array.rb
+++ b/mrbgems/mruby-array-ext/test/array.rb
@@ -409,3 +409,14 @@ assert("Array#combination") do
   assert_equal([[]], a.combination(0).to_a)
   assert_equal([], a.combination(5).to_a)
 end
+
+assert('Array#transpose') do
+  assert_equal([].transpose, [])
+  assert_equal([[]].transpose, [])
+  assert_equal([[1]].transpose, [[1]])
+  assert_equal([[1,2,3]].transpose, [[1], [2], [3]])
+  assert_equal([[1], [2], [3]].transpose, [[1,2,3]])
+  assert_equal([[1,2], [3,4], [5,6]].transpose, [[1,3,5], [2,4,6]])
+  assert_raise(TypeError) { [1].transpose }
+  assert_raise(IndexError) { [[1], [2,3,4]].transpose } 
+end