Created
September 11, 2011 02:11
-
-
Save cvrebert/1209079 to your computer and use it in GitHub Desktop.
Refactored "recursive algorithm for balls in numbered boxes"
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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