Home Zip like function that iterates over multiple items in lists and returns possibilities

Questions

Zip like function that iterates over multiple items in lists and returns possibilities

July 12, 2022

In the following code:

a = [["2022"], ["2023"]]
b = [["blue", "red"], ["green", "yellow"]]
c = [["1", "2", "3"], ["4", "5", "6", "7"], ["8", "9", "10", "11"], ["12", "13"]]

I would like a function that outputs this, but for any number of variables:

[
    ["2022", "blue", "1"],
    ["2022", "blue", "2"],
    ["2022", "blue", "3"],
    ["2022", "red", "4"],
    ["2022", "red", "5"],
    ["2022", "red", "6"],
    ["2022", "red", "7"],
    ["2023", "green", "8"],
    ["2023", "green", "9"],
    ["2023", "green", "10"],
    ["2023", "green", "11"],
    ["2023", "yellow", "12"],
    ["2023", "yellow", "13"],
]

I have searched for a function to do this with itertools or zip, but haven’t found anything yet.

>Solution :

First, you join the first argument, to a list of lists with only one element each.

Then for each sublist and its index i in the next argument, you pick the i-th list of the previous iteration res[i] and add to aux len(sublist) lists each of one is the res[i] with one item from sublist.

def f(*args):
    res = sum([[[item] for item in l] for l in args[0]], start=[])
    for arg in args[1:]:
        aux = []
        for i, sublist in enumerate(arg):
            aux += [res[i] + [opt] for opt in sublist]
        res = aux
    return res

In addition if you want to verify that the arguments passed to the function are correct, you can use this:

def check(*args):
    size = sum(len(l) for l in args[0])
    for arg in args[1:]:
        if len(arg) != size:
            return False
        size = sum(len(l) for l in arg)
    return True