Условное декартово произведение списков в itertools

Вот интерактивная демонстрация решения с использованием генератора.

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

Обратите внимание, что порядок отличается от того, который запрашивал Майкл, потому что второй компонент в 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')

Вы можете взять product, заменив два последних списка объектом range равной длины, а затем выбрать любой из двух последних списков в зависимости от последнего символа элемента из 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')

Используя 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')]

Хорошо, я попробовал. Итак, вы знаете первую часть:

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

Вот код, если вы хотите поиграть с ним.