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

And*_*Mao 11 python numpy permutation combinatorics scipy

假设我们有n投掷k球的垃圾箱.什么是快速(即使用numpy/scipy而不是python代码)方式来生成所有可能的结果作为矩阵?

例如,如果n = 4k = 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
Run Code Online (Sandbox Code Playgroud)

如果错过任何排列,请道歉,但这是一般的想法.生成的排列不必具有任何特定的顺序,但上述列表便于在心理上明确地迭代它们.

更好的是,有没有办法将每个从1的整数映射到多重编号(此列表的基数)直接映射到给定的排列?

这个问题与以下问题有关,这些问题在R中实现,具有非常不同的设施:

还有相关参考:

Jar*_*uen 3

这是一个使用 的生成器解决方案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]]
Run Code Online (Sandbox Code Playgroud)

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

请注意,虽然 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
Run Code Online (Sandbox Code Playgroud)