多个列表的 Itertools 组合，每个列表选择 n 个元素答案

【问题标题】：Itertools combinations of multiple list selecting n elements per list多个列表的 Itertools 组合，每个列表选择 n 个元素
【发布时间】：2021-11-14 20:49:16
【问题描述】：

我需要组合一个列表列表，例如从每个列表中选择 n 个元素

a=[[1,2,3,4,5],[6,7,8,9,10]]
n1=2
n2=3

所以我的结果可能是这样的：

r=[[1,2,6,7,8],[1,2,6,7,9],...,[4,5,7,8,9],[4,5,8,9,10]]

有什么干净的方法吗？还是我必须将列表分成更小的尺寸并使用 for 循环来调用 itertools？

【问题讨论】：

标签： python combinations itertools

【解决方案1】：

简单地分别生成两个列表的组合，然后取两个生成器的笛卡尔积：

from itertools import product, combinations

r_gen  = product(combinations(a[0], n1), combinations(a[1], n2))                             

r = (a + b for a, b in r_gen)

r 产生的前 10 个元素是

[(1, 2, 6, 7, 8),
 (1, 2, 6, 7, 9),
 (1, 2, 6, 7, 10),
 (1, 2, 6, 8, 9),
 (1, 2, 6, 8, 10),
 (1, 2, 6, 9, 10),
 (1, 2, 7, 8, 9),
 (1, 2, 7, 8, 10),
 (1, 2, 7, 9, 10),
 (1, 2, 8, 9, 10)]

【讨论】：

我想你在r = list(a + b for a, b in r_gen) 中想念list r 仍然是generator。
感谢 Brian 的回答，如果我还有第三个列表 a[2], n3，我可以结合相同的方法吗？例如使用结果列表 r，其中 n_r=列表的长度： r_gen = product(combinations(r, n_r]), combination(a[2], n3)) ？还是有直接的方法？
@Pulse9 itertools.product 接受任意数量的可迭代对象，因此您可以简化所有三个组合迭代器的传递。您还可以使用 product(*[combinations(lst, num) for lst, num in zip(a, n)] 之类的东西将其概括为任意数量的列表，其中 a 是您的列表列表，n 是每个列表的选择列表。
太完美了！非常感谢布赖恩
@Pulse9 * 表示argument unpacking。它允许您将可迭代的元素作为单独的参数传递给函数，因此foo(*[a, b, c]) 与编写foo(a, b, c) 具有相同的效果。另请参阅Pass a list to a function to act as multiple arguments。

【解决方案2】：

如果我理解正确的话，这个问题基本上有两个步骤：

从每组中选择 n 个项目。这可以使用itertools.combinations（或.permutations，根据您的需要完成：

a1 = itertools.combinations(a[0], n1)
a2 = itertools.combinations(a[1], n2)

找到这两个迭代的组合。这几乎是笛卡尔积所做的：

r = itertools.product(a1, a2)

为了使结果看起来完全符合您的要求，您可以使用列表推导来连接元组：

r = [list(s1 + s2) for s1, s2 in r

【讨论】：

我认为你的意思是在你的第一个块中写a2 = itertools.combinations(a[1], n2)。
你是对的。谢谢@Brian