itertools中列表的条件笛卡尔积

时间:2017-10-31 20:02:42

标签: python itertools

我有四个清单:

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

我想获得LISTALISTB的笛卡尔积,然后根据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如何实现第二部分,我不确定最佳方法。建议?

4 个答案:

答案 0 :(得分:2)

您可以使用产品将最后两个列表替换为长度相等的 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')

答案 1 :(得分:1)

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

答案 2 :(得分:1)

好的,我去了。所以你知道第一部分:

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

Here代码,如果你想玩它。

答案 3 :(得分:1)

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

>>> 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)的反转结果。 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')

另请注意,此方法允许LISTCLISTD具有不等长度。 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')