我有四个清单:
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)
获得第一部分,但我一直在查看itertools
如何实现第二部分,我不确定最佳方法。建议?
答案 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')
另请注意,此方法允许LISTC
和LISTD
具有不等长度。 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')