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

7

好的。我正在寻找最聪明、最紧凑的方法来实现这个功能。

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的数组,并在同一范围内生成所有组合。
可能存在一种简单的方法,也许是通过连接数组或以某种方式嵌套理解列表,但我的解决方案似乎过于复杂了。
很抱歉发了这么长的帖子。

3
itertools 模块有一个 product 函数,正好可以实现你要的功能。使用 from itertools import product; product(range(6), range(6), range(6))。无论哪个迭代器放在最后,它都是循环速度最快的。 - Elliot
5个回答

10
你可以使用 itertools.product,它只是一个嵌套迭代器的便捷函数。如果你想要多次重复相同的可迭代对象,它还有一个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)]

使用 map 可以获取 list 的列表:

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

我终于用[list(elem) for elem in product(*repeat(iterable, times))]写出了一行代码。 - madtyn

6
你可以使用 itertools.product
from itertools import product

def f(minimum, maximum, n):
    return list(product(*[range(minimum, maximum)] * n))

list替换为生成器,以提高内存效率。

2

itertools 拥有你所需的一切。 combinations_with_replacement 会从给定的可迭代对象中生成指定长度的、包含重复元素的组合。请注意,返回的值将是一个迭代器。

def f(min, max, num):    
    return itertools.combinations_with_replacement(range(min, max), num)

很好用于其他情况,但我需要[0,0,1] != [1,0,0]。顺序确实很重要。也许我应该解释得更清楚些。不过我会先把这个写下来的。 - madtyn
好的,正如其他答案所述 - itertools.product会更好。 - Pearley

1

一个纯Python实现:

k=2  # k-uples
xmin=2
xmax=5 
n=xmax-xmin

l1 = [x for x in range(n**k)]
l2 = [[ x//n**(k-j-1)%n for x in l1] for j in range(k)]          
l3 = [[ xmin + l2[i][j] for i in range(k)] for j in range(n**k)]

l3是:

[[2 2]
 [2 3]
 [2 4]
 [3 2]
 [3 3]
 [3 4]
 [4 2]
 [4 3]
 [4 4]]

你应该检查一下代码。*//会产生语法错误 - madtyn
更好的称呼应该是“纯内置”实现。我认为它不符合Pythonic的风格,因为它难以阅读,并且可能比使用Python标准库中的功能产生更多的内存和时间开销。 - Pearley

-1
你需要的是范围的笛卡尔积。幸运的是,这已经在itertools中存在了。
import itertools
print(list(itertools.product(range(0,5), range(0,5), range(0,5))))

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接