I have a number of lists as follows:
LISTA = ['A1', 'A2']
LISTB = ['B1', 'B2']
LISTC = ['C1']
LISTD = ['D1', 'D2']
LISTE = ['E1', 'E2', 'E3']
I'd like to write a function to get the Cartesian product of LISTX
and LISTY
, and then depending on the value of LISTY
in each tuple, I'd like to add either the product of LISTJ
, or the product of LISTK
. All as named key/value pairs, or as lists where [0]
is the key and [1]
is the value.
i.e. If X == LISTA
, Y == LISTB
, J == LISTC
, K == LISTE
, and we assume if there is value B1
we get the product with J(LISTC)
and B2
we get the product with K(LISTE)
, then we should get something like...
[{'LISTA': A1, 'LISTB': B1 'LISTC': C1},
{'LISTA': A1, 'LISTB': B2 'LISTE': E1},
{'LISTA': A1, 'LISTB': B2 'LISTE': E2},
{'LISTA': A1, 'LISTB': B2 'LISTE': E3},
{'LISTA': A2, 'LISTB': B1 'LISTC': C1},
{'LISTA': A2, 'LISTB': B2 'LISTE': E1},
{'LISTA': A2, 'LISTB': B2 'LISTE': E2},
{'LISTA': A2, 'LISTB': B2 'LISTE': E3}]
I think it's probably something along these lines, where we first get the product, and then use a generator in order to return a dict for each tuple in that product, along with the assigned k,v
pair from j
or k
. However, I can't figure out how to make this more generic than having 'y_val' in the arguments, and I'm also not sure how to assign the keys in J_OR_K
, rather than just ending up with a tuple like (('LISTA', A1), ('LISTB', B1), C1)
.
import itertools as it
def J_OR_K(P, y_val1, y_val2, j_list, k_list):
for k, v in P:
for x in {y_val1:j_list, y_val2:k_list}[v]:
yield {a,b,x}
def my_func(x, y, j, k):
res = []
# get cartesian product of x, y
prod_xy = it.product(it.product(['LISTA'], x),
it.product(['LISTB'], y))
# get product with EITHER J OR K
for t in J_OR_K(j, k):
res.append(t)
Although this is similar to a question I've previously asked, this is moderately different and more involved - so felt it required a new question.
Here is an interactive demonstration of a solution:
>>> import itertools
>>> LISTA = ['A1', 'A2']
>>> LISTB = ['B1', 'B2']
>>> LISTC = ['C1']
>>> LISTD = ['D1', 'D2']
>>> LISTE = ['E1', 'E2', 'E3']
>>> def f():
... for a,b in itertools.product(LISTA,LISTB):
... n, l = {"B1":["LISTC",LISTC], "B2":["LISTE",LISTE]}[b]
... for x in l:
... yield {"LISTA":a, "LISTB":b, n:x}
...
>>> from pprint import pprint
>>> pprint(list(f()))
[{'LISTA': 'A1', 'LISTB': 'B1', 'LISTC': 'C1'},
{'LISTA': 'A1', 'LISTB': 'B2', 'LISTE': 'E1'},
{'LISTA': 'A1', 'LISTB': 'B2', 'LISTE': 'E2'},
{'LISTA': 'A1', 'LISTB': 'B2', 'LISTE': 'E3'},
{'LISTA': 'A2', 'LISTB': 'B1', 'LISTC': 'C1'},
{'LISTA': 'A2', 'LISTB': 'B2', 'LISTE': 'E1'},
{'LISTA': 'A2', 'LISTB': 'B2', 'LISTE': 'E2'},
{'LISTA': 'A2', 'LISTB': 'B2', 'LISTE': 'E3'}]
>>>