python - 生成 n 个箱子中 k 个球的所有可能结果(多项式/分类结果的总和)

Question

假设我们有 n 个垃圾箱，我们要在其中扔 k 个球。什么是快速(即使用 numpy/scipy 而不是 python 代码)将所有可能的结果生成为矩阵的方法？

例如，如果 n = 4 和 k = 3，我们需要以下 numpy.array:

3 0 0 0
2 1 0 0
2 0 1 0
2 0 0 1
1 2 0 0
1 1 1 0
1 1 0 1
1 0 2 0
1 0 1 1
1 0 0 2
0 3 0 0
0 2 1 0
0 2 0 1
0 1 2 0
0 1 1 1
0 1 0 2
0 0 3 0
0 0 2 1
0 0 1 2
0 0 0 3

如果遗漏了任何排列，我们深表歉意，但这是一般的想法。生成的排列不必按任何特定顺序排列，但上面的列表很方便在脑海中对它们进行分类迭代。

更好的是，有没有一种方法可以将每个整数从 1 映射到 multiset number (此列表的基数)直接到给定的排列？

这个问题与以下问题有关，它们是在 R 中使用非常不同的设施实现的:

• Generating all permutations of N balls in M bins
• Generate a matrix of all possible outcomes for throwing n dice (ignoring order)

还有相关引用资料:

• https://en.wikipedia.org/wiki/Stars_and_bars_(combinatorics)
• https://en.wikipedia.org/wiki/Multiset#Counting_multisets
• https://en.wikipedia.org/wiki/Combinatorial_number_system

Answer 1

这里有一个使用 itertools.combinations_with_replacement 的生成器解决方案，不知道是否适合您的需求。

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]]

此函数的复杂性呈指数级增长，因此可行与不可行之间存在离散边界。

请注意，虽然 numpy 数组在构造时需要知道它们的大小，但这很容易实现，因为多重集数很容易找到。下面可能是更好的方法，我没有做过计时。

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