Enumeration of combinations of N balls in A boxes?

Question

I want to enumerate all possible combinations of N balls in A boxes.

example: I have 8

Answer 1

def iterate_assignments(N, K):
    buckets = [0] * K
    buckets[0] = K
    while True:
        yield buckets
        if buckets[-1] == N:
            return
        non_empty_buckets = sorted([i for i, count in enumerate(buckets) if count > 0])
        if non_empty_buckets[-1] == K - 1:
            temp = buckets[-1]
            buckets[-1] = 0
            buckets[non_empty_buckets[-2] + 1] = temp + 1
            buckets[non_empty_buckets[-2]] -= 1
        else:
            buckets[non_empty_buckets[-1]] -= 1
            buckets[non_empty_buckets[-1] + 1] += 1

This is a solution in python which efficiently yields all possible assignments of N balls to K buckets in O(K) memory and O(# possible assignments) time.

Start with an assignment where all balls are in the left-most bucket (for the purposes of understanding, assume all buckets are arranged left-to-right). Generate all later assignments by progressively moving balls to the right in the following fashion:

(1) if the right-most ball is in the right-most bucket, find the next right-most ball and move both of these balls into the buckets one to the right of the next right-most ball (2) If the right-most ball is anywhere else, move it one bucket to the right

From this scheme, its clear to see why this only generates unique combinations. It takes a bit more thought to see why this produces all unique combinations. I can attempt to formalize a proof if people are interested, but I'm skipping for now :)