Simple Ruby algorithm to generate a weighted list

weighted_list.rb

# README
#
# To run tests:
# $ ruby weighted_round_robin.rb
# PASS [[0, "a"]] makes []
# PASS [[1, "a"]] makes ["a"]
# PASS [[1, "a"], [2, "b"]] makes ["b", "a", "b"]
# ...
#
# To debug (prints index as ID):
# $ ruby 1 2 3
# 2 1 0 2 1 2
#
#
# DOCS
#
# Parameters:
# `weights` is a two-dimensional array where each item of the outer array
# is a 2-tuple containing (in order) the identifier of the item and its weight
#
# Returns:
# A flat array of identifiers occuring as many times as their weight,
# distributed evenly throughout the array
def weighted_list(weights)
  weights.flat_map do |id, weight|
    step = 1 / (weight.to_f + 1)

    weight.times.map do |placement|
      [
        step * (placement + 1),
        id
      ]
    end
  end.sort.map(&:last)
end

TESTS = {
  [["a", 0]] => [],
  [["a", 1]] => %w[a],
  [["a", 1], ["b", 2]] => %w[b a b],
  [["a", 1], ["b", 2], ["c", 4]] => %w[c b c a c b c],
  [["a", 1], ["b", 4]] => %w[b b a b b],
  [["a", 1], ["b", 3]] => %w[b a b b],
  [["b", 3], ["a", 1]] => %w[b a b b],
  [["a", 1], ["b", 0]] => %w[a],
  [["a", 2], ["b", 2], ["c", 2], ["d", 3]] => %w[d a b c d a b c d],

  # BAD CASES: we wish these were more balanced (non-sequential heavies) but
  # I can't think of a heuristic that doesn't break our simple algorithm
  [["a", 1], ["b", 1], ["c", 2]] => %w[c a c b],
  [["a", 1], ["b", 1], ["c", 1], ["d", 2]] => %w[d a b d c],
  [["a", 2], ["b", 2], ["c", 2], ["d", 4]] => %w[d a b d c d a b d c], #
}

if ARGV.any?
  puts weighted_list(ARGV.map(&:to_i).each_with_index).join(' ')
else
  TESTS.each do |weights, test|
    result = weighted_list(weights)
    if result == test
      puts "--- PASS #{weights} makes #{result}"
    else
      puts "!!! FAIL #{weights} makes #{result} not #{test}"
    end
  end
end