Generate all possible outcomes of k balls in n bins (sum of multinomial / categorical outcomes)

Question

Suppose we have n bins in which we are throwing k balls. What is a fast (i.e. using numpy/scipy instead of python code) way to gen

Answer 1

Here's a generator solution using itertools.combinations_with_replacement, don't know if it will be suitable for your needs.

def partitions(n, b):
    masks = numpy.identity(b, dtype=int)
    for c in itertools.combinations_with_replacement(masks, n): 
        yield sum(c)

output = numpy.array(list(partitions(3, 4)))
# [[3 0 0 0]
#  [2 1 0 0]
#  ...
#  [0 0 1 2]
#  [0 0 0 3]]

The complexity of this function grows exponentially, so there is a discrete boundary between what is feasible and what is not.

Note that while numpy arrays need to know their size at construction, this is easily possible since the multiset number is easily found. Below might be a better method, I have done no timings.

from math import factorial as fact
from itertools import combinations_with_replacement as cwr

nCr = lambda n, r: fact(n) / fact(n-r) / fact(r)

def partitions(n, b):
    partition_array = numpy.empty((nCr(n+b-1, b-1), b), dtype=int)
    masks = numpy.identity(b, dtype=int)
    for i, c in enumerate(cwr(masks, n)): 
        partition_array[i,:] = sum(c)
    return partition_array