Python:生成范围内值的所有n长度数组组合

Question

好的。我正在寻找最智能、更紧凑的方式来执行此功能

def f():
    [[a,b,c] for a in range(6) for b in range(6) for c in range(6)]

它应该像这样生成值 a、b、c 的所有组合:

[0,0,0]
[0,0,1]
[0,0,2]
...
[1,0,0]
[1,0,1]
...

等等……

但我希望它是灵活的，这样我就可以更改范围或可迭代对象，以及生成的数组的长度。范围是一件容易的事:

def f(min, max):
    [[a,b,c] for a in range(min,max) for b in range(min,max) for c in range(min,max)]

这对于 3 长度数组没问题，但我现在正在考虑制作 4 长度数组或 7 长度数组，并在同一范围内为它们生成所有组合。

它必须存在一种简单的方法，也许以某种方式连接数组或嵌套理解列表，但我的解决方案似乎太复杂了。

抱歉发了这么长的帖子。

Answer 1

您可以使用 itertools.product这只是嵌套迭代的一个便利函数。如果您想多次重复相同的 iterable，它还有一个 repeat 参数:

>>> from itertools import product

>>> amin = 0
>>> amax = 2
>>> list(product(range(amin, amax), repeat=3))
[(0, 0, 0), (0, 0, 1), (0, 1, 0),  (0, 1, 1),  (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]

要获取 list 的 list，您可以使用 map:

>>> list(map(list, product(range(amin, amax), repeat=3)))
[[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]]

但是 product 是一个迭代器，因此如果您只是对其进行迭代而不是将其强制转换为 list，它会非常高效。至少如果这在你的程序中是可能的。例如:

>>> for prod in product(range(amin, amax), repeat=3):
...     print(prod)  # one example
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(0, 1, 1)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(1, 1, 1)