I have four lists:
LISTA = ['A1', 'A2']
LISTB = ['B1_C', 'B2_D']
LISTC = ['C1', 'C2']
LISTD = ['D1', 'D2']
I'd like to get the Cartesian product of LISTA and LISTB, and then depending on the value of B, I'd like to add either the product of C, or the product of 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)
I can get the first part with itertools.product(LISTA, LISTB), but I've been looking through itertools for how to achieve the second part and I'm not sure the best way to go. Suggestions?
Here is an interactive demonstration of a solution using a generator.
>>> 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')
Note that the order is different then what Michael requested because the second component in product(LISTA,LISTB) changes faster then the first.
To get the exact ordering specified we need the reversed results from product(LISTB,LISTA). E.g.
>>> 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')
Note also that this approach allows LISTC and LISTD to have unequal length. E.g.
>>> 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')