Search code examples
pythonpython-itertoolsnested-lists

Conditional Cartesian product of a X lists, resulting in dict with k,v pairs?


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.


Solution

  • 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'}]
    >>>