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

For reference purposes, the following code uses Ehrlich's algorithm to iterate through all possible combinations of a multiset in C++, Javascript, and Python:

https://github.com/ekg/multichoose

This can be converted to the above format using this method. Specifically,

for s in multichoose(k, set):    row = np.bincount(s, minlength=len(set) + 1)

This still isn't pure numpy, but can be used to fill a preallocated numpy.array pretty quickly.

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 factfrom itertools import combinations_with_replacement as cwrnCr = 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

here is a naive implementation with list comprehensions, not sure about performance compared to numpy

def gen(n,k):    if(k==1):        return [[n]]    if(n==0):        return [[0]*k]    return [ g2 for x in range(n+1) for g2 in [ u+[n-x] for u in gen(x,k-1) ] ]> gen(3,4)[[0, 0, 0, 3], [0, 0, 1, 2], [0, 1, 0, 2], [1, 0, 0, 2], [0, 0, 2, 1], [0, 1, 1, 1], [1, 0, 1, 1], [0, 2, 0, 1], [1, 1, 0, 1], [2, 0, 0, 1], [0, 0, 3, 0], [0, 1, 2, 0], [1, 0, 2, 0], [0, 2, 1, 0], [1, 1, 1, 0], [2, 0, 1, 0], [0, 3, 0, 0], [1, 2, 0, 0], [2, 1, 0, 0], [3, 0, 0, 0]]