执行 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))。

但是现在假设我想找到n_A 重复A 和n_B 重复B 的笛卡尔积，其中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)