Skip to content

Instantly share code, notes, and snippets.

@replaid
Created October 22, 2016 21:09
Show Gist options
  • Save replaid/686049de192a4af65c267d1ea0110ed2 to your computer and use it in GitHub Desktop.
Save replaid/686049de192a4af65c267d1ea0110ed2 to your computer and use it in GitHub Desktop.
Flatten arbitrarily nested collections.
class FlattenEnum
class << self
# Returns an enumerator that yields the non-enumerable items from within a
# nested enumerable, in depth-first order.
#
# Example:
#
# FlattenEnum.create([[1, 2, [3]], 4]).to_a
# #=> [1, 2, 3, 4]
def create(nested_collection)
Enumerator.new do |yielder|
# The single-item case. Yield the item and stop the iteration.
unless nested_collection.respond_to?(:each)
item = nested_collection
yielder << item
raise StopIteration
end
# Recurse over each of the collection's elements in case they need
# flattening.
nested_collection.each do |element|
inner_enum = create(element)
loop do
yielder << inner_enum.next
end
end
end
end
end
end
require 'minitest/autorun'
require_relative 'flatten_enum'
class TestFlatten < Minitest::Test
def flatten(nested_collection)
enum = FlattenEnum.create(nested_collection)
enum.to_a
end
def test_single_integer
input = [1]
assert_equal([1], flatten(input))
end
def test_two_integers
input = [2, 3]
assert_equal([2, 3], flatten(input))
end
def test_single_nested_integer
input = [[4]]
assert_equal([4], flatten(input))
end
def test_full_customer_example
input = [[1, 2, [3]], 4]
expected = [1, 2, 3, 4]
assert_equal(expected, flatten(input))
end
def nest_pathologically(item)
[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[ item ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]
end
def test_pathological_depth
input = nest_pathologically(
[1, 2, nest_pathologically([3, nest_pathologically(4)])]
)
expected = [1, 2, 3, 4]
assert_equal(expected, flatten(input))
end
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment