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

user_19536130 0
Viewed 234

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.

To clarify, my use case for this was to iterate through values of a nested/multi-level dropdown menu (the first dropdown returns options, and each option returns a different dropdown, and so on).

python
4 Answers

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.

from itertools import chain

def f(*args):
    res = list(chain.from_iterable([[item] for item in l] for l in args[0]))
    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
user_194244910 
def foo(*args):
    arrs = list(args)
    prev = arrs[0]
    
    def bar(arr1, arr2):
        ans = []
        for x, y in zip(arr1, arr2):
            for el in y: 
                ans.append(x + [el])
        return ans
    
    for curr in arrs[1:]:
        ans = bar(prev, curr)
        prev = ans
    
    return ans

foo(a, b, c)
# [['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']]
user_71289340

Like d.b's answer, but using reduce and list comprehension:

from functools import reduce
from pprint import pprint

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

def foo(*args):
    def bar(arr1, arr2):
        return [
            x + [el]
            for x, y in zip(arr1, arr2)
            for el in y
        ]
    return reduce(bar, args)

pprint(foo(a, b, c))

Output (Try it online!):

[['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']]
user_126710570

If possible, I would use a different set of inputs of nested dicts/lists:

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

This makes the tree structure clearer, and would be safer since it's not possible that the list sizes mismatch. Though duplicate keys would no longer be possible.

We can the use recursion with dict defining branches and list defining leaf nodes.

def iter_options(options, path=()):
    if isinstance(options, dict):
        for key, value in options.items():
            yield from iter_options(value, path + (key,))
    elif isinstance(options, list):
        for key in options:
            yield path + (key,)
    else:
        raise TypeError

Usage:

>>> for row in iter_options(options):
...     print(row)
...
('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')
user_40009640
Related