itertools中列表的条件笛卡尔积

Question

我有四个清单：

LISTA = ['A1', 'A2']
LISTB = ['B1_C', 'B2_D']
LISTC = ['C1', 'C2']
LISTD = ['D1', 'D2']

我想获得LISTA和LISTB的笛卡尔积，然后根据B的值，我想添加C的乘积或D的乘积。

(A1 B1_C C1)
(A1 B1_C C2)
(A2 B1_C C1)
(A2 B1_C C2)
(A1 B2_D D1)
(A1 B2_D D2)
(A2 B2_D D1)
(A2 B2_D D2)

我可以使用itertools.product(LISTA, LISTB)来获得第一部分，但是我一直在通过itertools寻找如何实现第二部分，但我不确定最好的方法。 建议？

Answer 1

您可以将产品替换为具有相等长度的range对象的最后两个列表，然后根据LISTB项目的最后一个字符选择最后两个列表LISTB ：

from itertools import product

def func(lsta, lstb, lstc, lstd):
    for b, a, i in product(lstb, lsta, range(len(lstc))):
        yield a, b, lstc[i] if b.endswith('C') else lstd[i]

for tup in func(LISTA, LISTB, LISTC, LISTD):          
    print(tup)

---

('A1', 'B1_C', 'C1')
('A1', 'B1_C', 'C2')
('A2', 'B1_C', 'C1')
('A2', 'B1_C', 'C2')
('A1', 'B2_D', 'D1')
('A1', 'B2_D', 'D2')
('A2', 'B2_D', 'D1')
('A2', 'B2_D', 'D2')

Answer 2

我认为使用itertools可以完成以下工作：

import itertools

LISTA = ['A1', 'A2']
LISTB = ['B1_C', 'B2_D']
LISTC = ['C1', 'C2']
LISTD = ['D1', 'D2']
res = []

dictb = {b:b.split("_")[1] for b in LISTB}

def product_for(lst, b, otherlst, result):
    for el in itertools.product(*[lst , [b] , otherlst]):
      result.append(el)

for k,v in dictb.items():
  if v == 'C':
    product_for(LISTA, k, LISTC,res)
  else:
    product_for(LISTA, k, LISTD,res)

print(res)

=> [('A1', 'B1_C', 'C1'), ('A1', 'B1_C', 'C2'), ('A2', 'B1_C', 'C1'), ('A2', 'B1_C', 'C2'), ('A1', 'B2_D', 'D1'), ('A1', 'B2_D', 'D2'), ('A2', 'B2_D', 'D1'), ('A2', 'B2_D', 'D2')]

Answer 3

好吧，我努力了。 因此，您知道第一部分：

part1 = itertools.product(LISTA, LISTB)

结果是：

[('A1', 'B1_C'), ('A1', 'B2_D'), ('A2', 'B1_C'), ('A2', 'B2_D')]

然后，您可以按每个元组的最后一个元素的最后一个字符分组：

keyfunc = lambda x: x[1][-1:]
grouped = itertools.groupby(sorted(part1, key=keyfunc), keyfunc)    
# convert group object to dictionary
grouped_dict = dict((k, list(v)) for k, v in grouped)

这给你这个：

{'C': [('A1', 'B1_C'), ('A2', 'B1_C')], 'D': [('A1', 'B2_D'), ('A2', 'B2_D')]}

现在，您可以在每个组上做一个产品，然后将它们加入备份：

c = itertools.product(grouped_dict['C'], LISTC)
d = itertools.product(grouped_dict['D'], LISTD)    
part2 = itertools.chain(c, d)

剩下的就是：

[(('A1', 'B1_C'), 'C1'),
 (('A1', 'B1_C'), 'C2'),
 (('A2', 'B1_C'), 'C1'),
 (('A2', 'B1_C'), 'C2'),
 (('A1', 'B2_D'), 'D1'),
 (('A1', 'B2_D'), 'D2'),
 (('A2', 'B2_D'), 'D1'),
 (('A2', 'B2_D'), 'D2')]

最后，您可以再次展平每个元素：

part2 = itertools.imap(lambda x: x[0] + (x[1],), part2)

这将为您带来最终结果：

[('A1', 'B1_C', 'C1'),
 ('A1', 'B1_C', 'C2'),
 ('A2', 'B1_C', 'C1'),
 ('A2', 'B1_C', 'C2'),
 ('A1', 'B2_D', 'D1'),
 ('A1', 'B2_D', 'D2'),
 ('A2', 'B2_D', 'D1'),
 ('A2', 'B2_D', 'D2')]

如果您想使用它， 这里是代码。

Answer 4

这是使用生成器的解决方案的交互式演示。

>>> import itertools
>>> LISTA = ['A1', 'A2']
>>> LISTB = ['B1_C', 'B2_D']
>>> LISTC = ['C1', 'C2']
>>> LISTD = ['D1', 'D2']
>>> def C_OR_D(P):
...    for a,b in P:
...      for x in {"C":LISTC, "D":LISTD}[b[-1]]:
...         yield a,b,x
... 
>>> for t in C_OR_D(itertools.product(LISTA,LISTB)):
...    print t
... 
('A1', 'B1_C', 'C1')
('A1', 'B1_C', 'C2')
('A1', 'B2_D', 'D1')
('A1', 'B2_D', 'D2')
('A2', 'B1_C', 'C1')
('A2', 'B1_C', 'C2')
('A2', 'B2_D', 'D1')
('A2', 'B2_D', 'D2')

---

请注意，顺序与Michael要求的顺序不同，因为product(LISTA,LISTB)的第二个组件更改速度比第一个组件快。

为了获得指定的确切顺序，我们需要product(LISTB,LISTA)的相反结果。 例如

>>> for t in C_OR_D((a,b) for (b,a) in itertools.product(LISTB,LISTA)):
...    print t
... 
('A1', 'B1_C', 'C1')
('A1', 'B1_C', 'C2')
('A2', 'B1_C', 'C1')
('A2', 'B1_C', 'C2')
('A1', 'B2_D', 'D1')
('A1', 'B2_D', 'D2')
('A2', 'B2_D', 'D1')
('A2', 'B2_D', 'D2')

还要注意，这种方法允许LISTC和LISTD具有不相等的长度。 例如

>>> LISTD = ['D1', 'D2', 'D3']
>>> for t in C_OR_D((a,b) for (b,a) in itertools.product(LISTB,LISTA)):
...    print t
... 
('A1', 'B1_C', 'C1')
('A1', 'B1_C', 'C2')
('A2', 'B1_C', 'C1')
('A2', 'B1_C', 'C2')
('A1', 'B2_D', 'D1')
('A1', 'B2_D', 'D2')
('A1', 'B2_D', 'D3')
('A2', 'B2_D', 'D1')
('A2', 'B2_D', 'D2')
('A2', 'B2_D', 'D3')