35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
def solution(num_buns, num_required):
|
|
duplication_factor = num_buns - num_required + 1
|
|
result = [[] for bun in range(num_buns)]
|
|
key_id = 0
|
|
for dist in generate_key_distribution(num_buns, duplication_factor):
|
|
for bun in range(num_buns):
|
|
if dist[bun]:
|
|
result[bun] += [key_id]
|
|
key_id += 1
|
|
return result
|
|
|
|
def generate_key_distribution(num_buns, duplication):
|
|
if num_buns == duplication:
|
|
yield [True for i in range(num_buns)]
|
|
elif num_buns > 0:
|
|
if duplication > 0:
|
|
for dist in generate_key_distribution(num_buns - 1, duplication - 1):
|
|
yield [True] + dist
|
|
for dist in generate_key_distribution(num_buns - 1, duplication):
|
|
yield [False] + dist
|
|
|
|
|
|
tests = [
|
|
((3, 1), [[0], [0], [0]]),
|
|
((2, 2), [[0], [1]]),
|
|
((3, 2), [[0, 1], [0, 2], [1, 2]]),
|
|
((2, 1), [[0], [0]]),
|
|
((4, 4), [[0], [1], [2], [3]]),
|
|
((5, 3), [[0, 1, 2, 3, 4, 5], [0, 1, 2, 6, 7, 8], [0, 3, 4, 6, 7, 9], [1, 3, 5, 6, 8, 9], [2, 4, 5, 7, 8, 9]])
|
|
]
|
|
|
|
for i, o in tests:
|
|
result = solution(*i)
|
|
print(i, result == o, result, o)
|