flatten.rb

#!/usr/bin/env rspec

shared_examples "#flatten" do

  it "should flatten empty array" do
    expect( flatten([]) ).to eq []
  end

  it "should flatten single level array" do
    expect( flatten([1, 2, 3]) ).to eq [1, 2, 3]
  end

  it "should flatten nested array" do
    expect( flatten([[1, 2, 3]]) ).to eq [1, 2, 3]
  end

  it "should flatten multilevel array" do
    expect( flatten([1, [2, 3], [4, 5, [], [6]]]) ).to eq [1, 2, 3, 4, 5, 6]
  end

  it "should flatten multilevel array" do
    expect( flatten([1, [2, 3], [4, 5, [], [6]]]) ).to eq [1, 2, 3, 4, 5, 6]
  end

  it "should raise error on recursive array" do
    a = []
    a << [1, a]
    expect{ flatten(a) }.to raise_error(ArgumentError)
  end

end

describe "normal flatten" do

  it_should_behave_like "#flatten"

  def flatten(array)
    result = []
    stack = []
    visited = {array => true}.compare_by_identity
    idx = 0
    loop do

      if array.length < idx + 1
        return result if stack.empty?
        array, idx = stack.pop(2)
        visited.delete(array)
        next
      end

      el = array[idx]
      idx += 1

      if Array === el
        raise ArgumentError, "tried to flatten recursive array" if visited[el]
        visited[el] = true
        stack.push(array, idx)
        array, idx = el, 0
      else
        result << el
      end
    end
  end
end

describe "recursive flatten" do

  it_should_behave_like "#flatten"

  def flatten(array, result=[], visited={}.compare_by_identity)
    visited[array] = true
    array.each do |el|
      if Array === el
        raise ArgumentError, "tried to flatten recursive array" if visited[el]
        flatten(el, result, visited)
      else
        result << el
      end
    end
    visited.delete(array)
    result
  end
end