Interleave arrays in ruby

interleave.rb

# Given a set of arrays with potentially infinite lengths
# 1. interleave the elements from input collections into a
# final 1-dimensional array
# 2. ensure the interleaving follows a frequency, as example:
# - 1 element from collection A (frequency: 1)
# - 2 element from collection B (frequency: 2)
# - 3 element from collection C (frequency: 3)
# - 4 element from collection D (frequency: 4)

# helper to generate infinite enumerators
def col(str)
  Enumerator.new do |yielder|
    1.upto(Float::INFINITY) { |n| yielder << "#{str}#{n}" }
  end
end

def interleave(collections, frequencies, max_length)
  final = []
  collection_enums = collections.map(&:each)
  counts = Array.new(frequencies.length) { 0 }

  while final.length < max_length
    frequencies.each.with_index do |freq, idx|
      while counts[idx] < freq
        return final if final.length >= max_length

        begin
          final << collection_enums[idx].next
        rescue StopIteration
          counts[idx] = freq
          next
        end

        counts[idx] += 1
      end
    end
    counts = Array.new(frequencies.length) { 0 }
  end

  final
end

result = interleave([col('A'), col('B'), col('C'), col('D')], [1, 2, 3, 4], 10)
expected = %w[A1 B1 B2 C1 C2 C3 D1 D2 D3 D4]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

result = interleave([col('A'), col('B'), col('C'), col('D')], [1, 2, 3, 4], 20)
expected = %w[A1 B1 B2 C1 C2 C3 D1 D2 D3 D4 A2 B3 B4 C4 C5 C6 D5 D6 D7 D8]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

result = interleave([col('A'), col('B'), col('C'), col('D')], [1, 1, 1, 1], 10)
expected = %w[A1 B1 C1 D1 A2 B2 C2 D2 A3 B3]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

result = interleave([col('A'), col('B'), col('C'), col('D')], [4, 3, 2, 1], 10)
expected = %w[A1 A2 A3 A4 B1 B2 B3 C1 C2 D1]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

result = interleave([col('A'), col('B').take(2), col('C'), col('D')], [1, 1, 1, 1], 10)
expected = %w[A1 B1 C1 D1 A2 B2 C2 D2 A3 C3]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

result = interleave([col('A'), col('B').take(2), [], col('D')], [1, 1, 1, 1], 10)
expected = %w[A1 B1 D1 A2 B2 D2 A3 D3 A4 D4]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

####################

def interleave2(collections, frequencies, max_length)
  final = []
  collection_enums = collections.map(&:each)
  counts = frequencies.dup
  return final if counts.sum.zero?
  collections_depleted = collections.map(&:nil?)

  while final.length < max_length
    while counts.sum > 0
      collection_enums.each.with_index do |col, idx|
        next if counts[idx].zero?
        return final if final.length >= max_length

        begin
          final << col.next
        rescue StopIteration
          collections_depleted[idx] = true
          counts[idx] = 0
          next
        end

        counts[idx] -= 1
      end

      return final if collections_depleted.all? { |cd| cd == true }
    end
    counts = frequencies.dup
  end

  final
end

result = interleave2([col('a'), col('b'), col('c'), col('d')], [1, 4, 5, 2], 10)
expected = %w[a1 b1 c1 d1 b2 c2 d2 b3 c3 b4]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

# explain_frequencies([1, 4, 5, 2])
result = interleave2([col('a'), col('b'), col('c'), col('d')], [1, 4, 5, 2], 22)
expected = %w[a1 b1 c1 d1 b2 c2 d2 b3 c3 b4 c4 c5 a2 b5 c6 d3 b6 c7 d4 b7 c8 b8]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

result = interleave2([[], [], [], []], [1, 4, 5, 2], 24)
expected = %w[]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

result = interleave2([col('a').take(1), col('b'), col('c'), col('d')], [1, 4, 5, 2], 22)
expected = %w[a1 b1 c1 d1 b2 c2 d2 b3 c3 b4 c4 c5 b5 c6 d3 b6 c7 d4 b7 c8 b8 c9]
puts 'Expected: ' + expected.join(' ')
puts '     Got: ' + result.join(' ')
puts expected == result ? '✅' : '❌'

puts 'Big sort'
puts Benchmark.measure {
  result = interleave2([col('a'), col('b'), col('c'), col('d')], [1, 4, 3, 2], 100_000)
  puts '  Total: ' + result.length.to_s
  puts '   Head: ' + result.take(10).join(' ')
  puts ' Middle: ' + result[(result.length/2 - 5)..(result.length/2 + 5)].join(' ')
  puts '   Tail: ' + result.last(10).join(' ')
}

puts 'Huge sort'
puts Benchmark.measure {
  result = interleave2([col('a'), col('b'), col('c'), col('d')], [1, 4, 3, 2], 1_000_000)
  puts '  Total: ' + result.length.to_s
  puts '   Head: ' + result.take(10).join(' ')
  puts ' Middle: ' + result[(result.length/2 - 5)..(result.length/2 + 5)].join(' ')
  puts '   Tail: ' + result.last(10).join(' ')
}