Tuesday, September 11, 2007

Ruby: Array#chunk

Yesterday I was working on some code that needed to split an array into 3 different chunks. My pair, Philippe Hanrigou, wrote tests similar to the following.

unit_tests do
test "chunk evenly" do
assert_equal [[1], [2], [3]], [1, 2, 3].chunk(3)
end

test "divide evenly" do
assert_equal [[1], [2], [3]], [1, 2, 3] / 3
end

test "chunk unevenly" do
assert_equal [[1, 3], [2]], [1, 2, 3].chunk(2)
end

test "chunk creates empty array when there aren't enough elements in the array" do
assert_equal [[1], [], []], [1].chunk(3)
end
end

To make the above tests pass, I wrote the following code.

class Array
def chunk(number_of_chunks)
chunks = (1..number_of_chunks).collect { [] }
while self.any?
chunks.each do |a_chunk|
a_chunk << self.shift if self.any?
end
end
chunks
end
alias / chunk
end

The above code works, but I'm curious to see if there is a better solution. How would you make the tests pass?

Here's the code in it's entirety.

require 'rubygems'
require 'test/unit'
require 'dust'

class Array
def chunk(number_of_chunks)
chunks = (1..number_of_chunks).collect { [] }
while self.any?
chunks.each do |a_chunk|
a_chunk << self.shift if self.any?
end
end
chunks
end
alias / chunk
end

unit_tests do
test "chunk evenly" do
assert_equal [[1], [2], [3]], [1, 2, 3].chunk(3)
end

test "divide evenly" do
assert_equal [[1], [2], [3]], [1, 2, 3] / 3
end

test "chunk unevenly" do
assert_equal [[1, 3], [2]], [1, 2, 3].chunk(2)
end

test "chunk creates empty array when there aren't enough elements in the array" do
assert_equal [[1], [], []], [1].chunk(3)
end
end
Post a Comment