Flatten an Array

flatten.rb

require 'minitest/autorun'

def count_all_elements(array)
  element_count = 0

  array.each do |el|
    # Given Ruby's duck-typing, it makes more sense to check for iterable
    # objects than for integer-like objects. Otherwise we'd do more
    # thorough type-checking.
    if el.respond_to? :each
      element_count += count_all_elements(el)
    else
      element_count += 1
    end
  end

  element_count
end

def squish(old_array, new_array=nil, index=nil)
  if not new_array
    new_array_length = count_all_elements(old_array)

    # While Ruby has dynamic arrays, creating one of exactly the size we need
    # spares us all the automatic resizing that would happen behind the scenes.
    new_array = new_array || Array.new(new_array_length)
  end

  index = index || 0

  old_array.each do |el|
    if el.respond_to? :each
      squish(el, new_array, index)
      # There is probably a way to avoid recounting these sub-arrays, but then
      # that is a relatively cheap operation
      index += count_all_elements(el) - 1
    else
      new_array[index] = el
    end

    index += 1
  end

  new_array
end

describe '#squish' do
  it 'returns an unchanged array if there were no nestled arrays to flatten' do
    squish([1,2,3,4]).must_equal([1,2,3,4])
  end

  it 'flattens an array with one nestled array' do
    squish([1,2,3,[4],5]).must_equal([1,2,3,4,5])
  end

  it 'flattens an array with two nestled arrays' do
    squish([1,2,3,[4],5,[6]]).must_equal([1,2,3,4,5,6])
  end

  it 'flattens an array with a doubly nestled array with one element' do
    squish([1,2,3,[[4]],5,6]).must_equal([1,2,3,4,5,6])
  end

  it 'flattens an array with a doubly nestled array with two elements' do
    squish([1,2,3,[[4,5]],6]).must_equal([1,2,3,4,5,6])
  end

  it 'flattens an array with a triply nestled array' do
    squish([1,2,3,[[[4,5]],6]]).must_equal([1,2,3,4,5,6])
  end

  it 'flattens an array with two triply nestled arrays' do
    squish([1,[[[2]]],3,[[[4,5]],6]]).must_equal([1,2,3,4,5,6])
  end

  it 'flattens an array with an empty nestled array' do
    squish([1,2,3,4,[],5]).must_equal([1,2,3,4,5])
  end

  it 'flattens an empty array without an error' do
    squish([]).must_equal([])
  end

  # We would add tests for typechecking here if it were necessary in Ruby
end