多個列表的 Itertools 組合，每個列表選擇 n 個元素

Question

我需要組合一個列表列表，例如選擇每個列表的 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？

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

Answer 2

如果我理解正確的話，這個問題基本上有兩個步驟：

1. 從每組中選擇 n 個項目。 這可以通過itertools.combinations （或.permutations根據您的需要：

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

2. 找到這兩個可迭代對象的組合。 這幾乎就是笛卡爾積所做的：

r = itertools.product(a1, a2)

為了使結果看起來正是您要查找的內容，您可以使用列表理解來連接元組：

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