Conditional Cartesian product of lists in itertools

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

You can take the product replacing the last two lists with a range object of equal length, and then select either of the last two lists depending on the last character of the item from 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')

Using itertools, I think this should do the work:

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

Ok, I had a go at it. So you know the first part:

part1 = itertools.product(LISTA, LISTB)

Which results in:

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

You could then group by the last character of the last element of each tuple:

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)

That gives you this:

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

Now you can do a product on each group, and join them back up:

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

Which leaves you with:

[(('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')]

Finally, you can flatten each element again:

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

Which gets you the final result:

[('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')]

Here's the code if you want to play with it.

来源：https://stackoverflow.com/questions/47043660/conditional-cartesian-product-of-lists-in-itertools

标签

python

itertools