python - itertools中列表的条件笛卡尔积

我有四个清单：

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)得到第一部分，但是我一直在寻找如何实现第二部分，我不确定最好的方法。建议？

最佳答案

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

>>> 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')