pythonfor-loopnestedon-the-fly

How to generate nested for loops on the fly in python


Imagine having a list (d) that is made of other lists with different lengths. We want to create a new list that looks like the result bellow. But obviously in this case everything is clear, we have a, b and c so we can create nested loops.

a = [1, 2]
b = [10, 20]
c = [100, 200, 300]

d = [a, b, c]

ret = []
for a1 in a:
    for b1 in b:
        for c1 in c:
            ret.append([a1, b1, c1])

print ret

Which results in this:

[[1, 10, 100], [1, 10, 200], [1, 10, 300], [1, 20, 100], [1, 20, 200], [1, 20, 300], [2, 10, 100], [2, 10, 200], [2, 10, 300], [2, 20, 100], [2, 20, 200], [2, 20, 300]]

What if we only have d with unknown number of lists in which there are unknown number of elements like bellow?

d = [[1, 2, 3], [10], [100, 200, 300, 400]]

Solution

  • What you want is itertools.product:

    import itertools
    a = [1, 2]
    b = [10, 20]
    c = [100, 200, 300]
    d = [a, b, c]
    ret = list(itertools.product(a, b, c))  # Or list(itertools.product(*d))