README.md

Implementation of flatten method of Array#flatten

Requires Ruby 1.9.2 or above. This implementation adds Array#my_flatten to an Array object, which is itself a replica of Array#flatten which is done in C.

Run test

Make sure you have minitest gem installed: $ gem install minitest. Then run:

ruby my_flatten_test.rb

my_flatten.rb

class Array
  def my_flatten(level = nil)
    rb_flatten(self, [], level)
  end

  private
  # apply recursion based on the level
  # when no level provided, then produce a complete flatten array
  # when level is given, then produce a flatten array flattened till that certain level
  def rb_flatten(array, result, level)
    array.each do |value|
      if ((value.is_a? Array) && (level.nil? || (level && level > 0)))
        rb_flatten(value, result, (level.nil? ? level : ((level || 0 ) - 1)))
      else
        result << value
      end
    end
    return result
  end
end

# arr = [[1,[2, [55, 6]],[3]],4]

# p arr.my_flatten    #=> [1, 2, 55, 6, 3, 4]
# p arr.my_flatten(1) #=> [1, [2, [55, 6]], [3], 4]
# p arr.my_flatten(2) #=> [1, 2, [55, 6], 3, 4]
# p arr               #=> [[1,[2, [55, 6]],[3]],4]

my_flatten_test.rb

require File.expand_path('../array_flatten', __FILE__)
require 'minitest/autorun'

class ArrayFlattenTest < MiniTest::Unit::TestCase
  def setup
    @arr = [[1,[2, [55, 6]],[3]],4]
  end

  def test_my_flatten
    flatted_array = @arr.flatten
    my_flatted_array = @arr.my_flatten
    assert_equal flatted_array, my_flatted_array
  end

  def test_my_flatten_with_one_level
    flatted_array = @arr.flatten(1)
    my_flatted_array = @arr.my_flatten(1)
    
    assert_equal flatted_array, my_flatted_array
  end

  def test_my_flatten_with_two_level
    flatted_array = @arr.flatten(2)
    my_flatted_array = @arr.my_flatten(2)
    
    assert_equal flatted_array, my_flatted_array
  end
end