custom_flatten.rb

# test utility for the script
def assert(actual, expected)
  if actual == expected
    puts "ok #{expected}"
  else
    puts "ko, you expected: #{expected} but got #{actual}"
  end
end

# first implementation, more functional programming friendly, pure method
def custom_flatten(arg)
  if arg.respond_to?(:each)
    arg.inject([]) {|ar, elt| ar + flatten(elt) }
  else
    [arg]
  end
end

# benefit: only create one array, impure method though :)
def alternative_flatten(arg, result = [])
  if arg.respond_to?(:each)
    arg.each {|elt| flatten(elt, result) }
    result
  else
    result.push arg
  end
end

# tests
assert(custom_flatten(1),
                     [1])
assert(custom_flatten([[1,2,[3]],4]),
                      [1,2,3,4])
assert(custom_flatten([[ 1, 2, 3 ],[ 4, 5, 6, [7, 8] ],  9,  10]),
                      [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
# bonus
assert(custom_flatten([['a','b',['c']],'d']),
                      ['a','b','c','d'])