array_flattener.rb

require 'minitest/autorun'

##
# flatten takes an array of arbitrarily nested arrays of integers and returns a flattened array.
# Alternatively, if a block is given, the flattened version of the items is yielded
#
# :call-seq:
#   flatten([[1,2,[3]],4]) -> [1,2,3,4]
#
#   flatten([[1,2,[3]],4], &block) yields 1, 2, 3, 4

def flatten(arr)
  if block_given?
    if arr.is_a?(Array)
      arr.each do |item|
        flatten(item, &Proc.new)
      end
    else
      yield arr
    end
    return
  end

  res = []

  flatten(arr) do |item|
    res << item
  end

  res
end

class ArrayFlattenerTest < Minitest::Test

  def test_empty
    assert_equal [], flatten([])
  end

  def test_already_flat
    assert_equal [1,2,3], flatten([1,2,3])
  end

  def test_simple
    assert_equal [1,2,3], flatten([1,[2,3]])
  end

  def test_sample
    assert_equal [1,2,3,4], flatten([[1,2,[3]],4])
  end

  def test_complex
    arr = [
      [1],
      [2, [3,4]],
      [[[5,6,[7]]]]
    ]

    assert_equal [1,2,3,4,5,6,7], flatten(arr)
  end
end