python - 执行 itertools.product 允许不同迭代次数的重复

Question

假设我有两个可迭代对象 A = 'ab' 和 B = '12'。

itertools.product返回一个迭代器，它将生成 A 和 B 的笛卡尔积；例如

>>> list(itertools.product(A,B)) 
[('a', '1'), ('a', '2'), ('b', '1'), ('b', '2')].

该函数有一个可选的关键字参数repeat，可用于查找可迭代对象与其自身的笛卡尔积；例如

>>> list(itertools.product(A,repeat=2))
[('a', 'a'), ('a', 'b'), ('b', 'a'), ('b', 'b')]

并且相当于list(itertools.product(A,A))。

然后使用 repeat=2 和 A 和 B 给出

>>> list(itertools.product(A,B,repeat=2))
[('a', '1', 'a', '1'), ('a', '1', 'a', '2'), ('a', '1', 'b', '1'), ('a', '1', 'b', '2'), ('a', '2', 'a', '1'), ('a', '2', 'a', '2'), ('a', '2', 'b', '1'), ('a', '2', 'b', '2'), ('b', '1', 'a', '1'), ('b', '1', 'a', '2'), ('b', '1', 'b', '1'), ('b', '1', 'b', '2'), ('b', '2', 'a', '1'), ('b', '2', 'a', '2'), ('b', '2', 'b', '1'), ('b', '2', 'b', '2')]

并且相当于list(itertools.product(A,B,A,B))。

但现在假设我想找到 A 的 n_A 重复和 B 的 n_B 重复的笛卡尔积code>，其中 n_A 和 n_B 不必相同。我怎样才能做到这一点？如果repeat采用元组(n_A, n_B)并且我可以写

，那就太好了

list(itertools.product(A,B,repeat=(n_A,n_B)))

例如

list(itertools.product(A,B,repeat=(2,3))) == list(itertools.product(A,A,B,B,B))

但这似乎是不允许的。

请注意，实际上 (A,A,B,B,B) 与 (A,B,A,B,B) 是不同的产品，但是我无论如何都会对输出进行排序，所以我不关心输入的顺序。

Answer 1

使用tee “复制”每个可迭代对象，然后将它们展平为单个参数列表，并使用 * 解包将它们作为单独的参数传递给产品

from itertools import product, chain, tee

def myproduct(*iterables, repeat=1):
    if isinstance(repeat, int):
        return product(iterables, repeat)
    assert isinstance(repeat, tuple)
    args = chain(*map(tee, iterables, repeat))
    return product(*args)


A = 'ab'
B = '12'

n_A = 2
n_B = 3


result = list(product(A, A, B, B, B))
result2 = list(myproduct(A, B, repeat=(n_A, n_B)))

print(result == result2)