[Python]: 生成所有可能组合的数组

Question

我有一个非常简单的组合问题。我有两个数组 (a 和 b)。数组 a 表示数组 b 中三个插槽中的一个可以取值的所有值。数组 b 中每个插槽的值都可以在 1 到 5 之间。一个例子是 [1,4,5]。我想生成一个包含所有可能组合的数组 c。我想扩展基本示例以用于更大的数组。

输入：

a = [1, 2, 3, 4, 5]
b = [1, 2, 3]

输出：

c = [[1, 1, 1], [1, 1, 2],[1, 1, 3], [1, 1, 4], [1, 1, 5],
     [1, 2, 1], [1, 2, 2],[1, 2, 3], [1, 2, 4], [1, 2, 5],
     [1, 3, 1], [1, 3, 2],[1, 3, 3], [1, 3, 4], [1, 3, 5],
     [1, 4, 1], [1, 4, 2],[1, 4, 3], [1, 4, 4], [1, 4, 5],
     [1, 5, 1], [1, 5, 2],[1, 5, 3], [1, 5, 4], [1, 5, 5],
     [2, 1, 1], [2, 1, 2],[2, 1, 3], [2, 1, 4], [2, 1, 5],
     [2, 2, 1], [2, 2, 2],[2, 2, 3], [2, 2, 4], [2, 2, 5],
     [2, 3, 1], [2, 3, 2],[2, 3, 3], [2, 3, 4], [2, 3, 5],
     [2, 4, 1], [2, 4, 2],[2, 4, 3], [2, 4, 4], [2, 4, 5],
     [2, 5, 1], [2, 5, 2],[2, 5, 3], [2, 5, 4], [2, 5, 5],
     [3, 1, 1], [3, 1, 2],[3, 1, 3], [3, 1, 4], [3, 1, 5],
     [3, 2, 1], [3, 2, 2],[3, 2, 3], [3, 2, 4], [3, 2, 5],
     [3, 3, 1], [3, 3, 2],[3, 3, 3], [3, 3, 4], [3, 3, 5],
     [3, 4, 1], [3, 4, 2],[3, 4, 3], [3, 4, 4], [3, 4, 5],
     [3, 5, 1], [3, 5, 2],[3, 5, 3], [3, 5, 4], [3, 5, 5],
     [4, 1, 1], [4, 1, 2],[4, 1, 3], [4, 1, 4], [4, 1, 5],
     [4, 2, 1], [4, 2, 2],[4, 2, 3], [4, 2, 4], [4, 2, 5],
     [4, 3, 1], [4, 3, 2],[4, 3, 3], [4, 3, 4], [4, 3, 5],
     [4, 4, 1], [4, 4, 2],[4, 4, 3], [4, 4, 4], [4, 4, 5],
     [5, 5, 1], [5, 5, 2],[5, 5, 3], [5, 5, 4], [5, 5, 5],
     [5, 1, 1], [5, 1, 2],[5, 1, 3], [5, 1, 4], [5, 1, 5],
     [5, 2, 1], [5, 2, 2],[5, 2, 3], [5, 2, 4], [5, 2, 5],
     [5, 3, 1], [5, 3, 2],[5, 3, 3], [5, 3, 4], [5, 3, 5],
     [5, 4, 1], [5, 4, 2],[5, 4, 3], [5, 4, 4], [5, 4, 5],
     [5, 5, 1], [5, 5, 2],[5, 5, 3], [5, 5, 4], [5, 5, 5]]

上述问题的解决方案：

d = []
for i in range(len(a)):
    for j in range(len(a)):
        for k in range(len(a)):
            e = []
            e.append(i+1)
            e.append(j+1)
            e.append(k+1)
            d.append(e)

我希望能够找到一种更通用的方法，可以容纳更大的数组（见下文），而不需要使用嵌套的for循环结构。我在stackoverflow上搜索了一个可比较的例子，但未能找到。

输入：

a = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
b = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]

- nkz

---

看一下 Python 的 itertools。组合方法可以实现你想要的功能。 - dlshriver

从您的描述中，我认为您不需要一个数组来存储b，只需要一个整数值来指定数组的长度。 - wilkesybear

当我对一些公差变量运行蒙特卡罗模拟时，我也曾有过类似的问题，但我发现生成所有可能组合的数组并没有意义，因为很容易（参见您自己的代码）在运行时生成这些组合。 - roadrunner66

这对我来说更像笛卡尔积，而不是组合，如果是这样的话，它就是许多以前问题的重复。 - DSM

@DSM 你说得对，应该是“product”。 - dlshriver

Answer 1

你需要使用 itertools.product() 函数。

a = [1, 2, 3, 4, 5]
b = 3  # Actually, you just need the length of the array, values do not matter

c = itertools.product(a, repeat=b)

请注意，这将返回一个迭代器，您可能需要使用list()进行强制类型转换，但请注意，如果大小增长，这可能需要很长时间，并且会高度消耗内存。

- Delgan

Answer 2

在一般情况下，当然应该使用itertools模块，而在这种特殊情况下，可以使用itertools.product，如其他答案所述。

如果您想自己实现该函数，可以使用递归使其适用于任何数组大小。此外，您可能应该将其作为生成器函数（使用yield而不是return）来创建，因为结果可能相当长。您可以尝试类似以下的代码：

def combinations(lst, num):
    if num > 0:
        for x in lst:
            for comb in combinations(lst, num - 1):
                yield [x] + comb
    else:
        yield []

- tobias_k