module Itertools
def Itertools.combination_iter(items, value_range)
# Example:
#
# items = ["a", "b"], value_range = [1, 2]
# combinations:
# [{"a":1, "b":1}, {"a":2, "b":1}, {"a":1, "b":2}, {"a":2, "b":2}]
combination = Hash[items.zip([value_range[0]] * items.length)]
yield combination
has_next_combination = true
while has_next_combination
next_combination = false
index = 0
while index < items.length and not next_combination
cur_value = combination[items[index]] + 1
if cur_value > value_range[1] then
combination[items[index]] = value_range[0]
index += 1
elsif
combination[items[index]] = cur_value
next_combination = true
yield combination
end
end
has_next_combination = next_combination
end
end
def Itertools.combination_with_total_sum(items, total)
# Example:
#
# items = ["a", "b"], total_sum = 3
# combinations:
# [{"a":0, "b":3}, {"a":1, "b":2}, {"a":2, "b":1}, {"a":3, "b":0}]
if items.length == 1 then
return [{items[0] => total}]
else
result = []
(0..total).each do |i|
combination = {items[0] => i}
sublist = combination_with_total_sum(items[1..-1], total-i)
sublist.each do |item|
combination.update(item)
result.push(combination.clone)
end
end
return result
end
end
end
__END__
def test
input_items = ["a", "b", "c"]
input_range = [1, 3]
output = [{"a"=>1, "b"=>1, "c"=>1},
{"a"=>2, "b"=>1, "c"=>1},
{"a"=>3, "b"=>1, "c"=>1},
{"a"=>1, "b"=>2, "c"=>1},
{"a"=>2, "b"=>2, "c"=>1},
{"a"=>3, "b"=>2, "c"=>1},
{"a"=>1, "b"=>3, "c"=>1},
{"a"=>2, "b"=>3, "c"=>1},
{"a"=>3, "b"=>3, "c"=>1},
{"a"=>1, "b"=>1, "c"=>2},
{"a"=>2, "b"=>1, "c"=>2},
{"a"=>3, "b"=>1, "c"=>2},
{"a"=>1, "b"=>2, "c"=>2},
{"a"=>2, "b"=>2, "c"=>2},
{"a"=>3, "b"=>2, "c"=>2},
{"a"=>1, "b"=>3, "c"=>2},
{"a"=>2, "b"=>3, "c"=>2},
{"a"=>3, "b"=>3, "c"=>2},
{"a"=>1, "b"=>1, "c"=>3},
{"a"=>2, "b"=>1, "c"=>3},
{"a"=>3, "b"=>1, "c"=>3},
{"a"=>1, "b"=>2, "c"=>3},
{"a"=>2, "b"=>2, "c"=>3},
{"a"=>3, "b"=>2, "c"=>3},
{"a"=>1, "b"=>3, "c"=>3},
{"a"=>2, "b"=>3, "c"=>3},
{"a"=>3, "b"=>3, "c"=>3}]
result = []
Itertools.combination_iter(["a", "b", "c"], [1, 3]) do |combination|
result.push(combination.clone)
end
puts "Test combination_iter: #{output == result ? 'passed' : 'failed'}"
input_items = ["a", "b", "c"]
input_total = 5
output = [{"a"=>0, "b"=>0, "c"=>5},
{"a"=>0, "b"=>1, "c"=>4},
{"a"=>0, "b"=>2, "c"=>3},
{"a"=>0, "b"=>3, "c"=>2},
{"a"=>0, "b"=>4, "c"=>1},
{"a"=>0, "b"=>5, "c"=>0},
{"a"=>1, "b"=>0, "c"=>4},
{"a"=>1, "b"=>1, "c"=>3},
{"a"=>1, "b"=>2, "c"=>2},
{"a"=>1, "b"=>3, "c"=>1},
{"a"=>1, "b"=>4, "c"=>0},
{"a"=>2, "b"=>0, "c"=>3},
{"a"=>2, "b"=>1, "c"=>2},
{"a"=>2, "b"=>2, "c"=>1},
{"a"=>2, "b"=>3, "c"=>0},
{"a"=>3, "b"=>0, "c"=>2},
{"a"=>3, "b"=>1, "c"=>1},
{"a"=>3, "b"=>2, "c"=>0},
{"a"=>4, "b"=>0, "c"=>1},
{"a"=>4, "b"=>1, "c"=>0},
{"a"=>5, "b"=>0, "c"=>0}]
result = Itertools.combination_with_total_sum(["a", "b", "c"], 5)
puts "Test combination_with_total_sum: #{output == result ? 'passed' : 'failed'}"
end
test