[英]Conditional Cartesian product of lists in itertools
我有四個清單:
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
尋找如何實現第二部分,但我不確定最好的方法。 建議?
您可以將產品替換為具有相等長度的range
對象的最后兩個列表,然后根據LISTB
項目的最后一個字符選擇最后兩個列表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')]
如果您想使用它, 這里是代碼。
這是使用生成器的解決方案的交互式演示。
>>> 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)
的相反結果。 例如
>>> 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')
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.