Refactored "recursive algorithm for balls in numbered boxes"
def balls_in_numbered_boxes(balls, box_sizes): | |
if not isinstance(balls, int): | |
raise TypeError("balls must be a non-negative integer.") | |
if balls < 0: | |
raise ValueError("balls must be a non-negative integer.") | |
box_sizes = list(box_sizes) | |
if not box_sizes: | |
raise ValueError("box_sizes must be a non-empty iterable.") | |
capacity = 0 | |
# capacity = sum(box_sizes) | |
for size in box_sizes: | |
if not isinstance(size, int): | |
raise TypeError("box_sizes must contain only positive integers.") | |
if size < 1: | |
raise ValueError("box_sizes must contain only positive integers.") | |
capacity += size | |
if capacity < balls: | |
raise ValueError("The total capacity of the boxes is less than the " | |
"number of balls to be distributed.") | |
return _balls_in_numbered_boxes(balls, box_sizes) | |
def _balls_in_numbered_boxes(balls, box_sizes): | |
if not balls: | |
yield len(box_sizes) * (0,) | |
elif len(box_sizes) == 1: | |
if box_sizes[0] >= balls: | |
yield (balls,) | |
else: | |
balls_in_first_box = min(balls, box_sizes[0]) | |
balls_in_other_boxes = balls - balls_in_first_box | |
rest_box_sizes = box_sizes[1:] | |
while True: | |
remaining_combos = _balls_in_numbered_boxes(balls_in_other_boxes, | |
rest_box_sizes) | |
for combo in remaining_combos: | |
yield (balls_in_first_box,) + combo | |
if not balls_in_first_box: | |
break | |
balls_in_first_box -= 1 | |
balls_in_other_boxes += 1 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment