Function with varying number of for loops using the arguments in the loop

Question

Firstly I would like to say that I've read this thread:

Function with varying number of For Loops (python)

And while it does certainly look similar to what I'd like to know, I'm not proficient enough (at all) with recursion to make it work.

If I have the following snippet:

I = ["i" + str(i+1) for i in range(3)]
J = ["j" + str(i+1) for i in range(4)]
K = ["k" + str(i+1) for i in range(2)]

Then I can make a function that does something like this:

def f(*Sets):
    size = len(Sets)
    if size == 1:
        for i in Sets[0]:
            print(i)
    elif size == 2:
        for i in Sets[0]:
            for j in Sets[1]:
                print(i,j)
    elif size == 3:
        for i in Sets[0]:
            for j in Sets[1]:
                for k in Sets[2]:
                    print(i,j,k)

And when executing, it would produce something like:

>>f(I)
i1
i2
i3

>>f(I,J)
i1,j1
i1,j2
...
i3,j4

>>f(I,J,K)
i1,j1,k1
i1,j1,k2
...
i3,j4,k2

The print statement is just an example. I would like to actually access elements in a pandas dataframe at the same time as checking the looped variable, or a dictionary, so it would include some command like print(i,j,k, ":", dictionary[i,j,k]). I need to access each variable in the loops, that´s what I mean.

This looks to me that could be widely reduced using some kind of recursion. Because if for example I need four lists, I'd have to add another elif. But I have not a recursive type of mind, and can't see how to solve it.

Or maybe it cannot be done at all. Not sure anymore of anything :P

user2390182 · Accepted Answer · 2020-02-14T09:47:58.857

5

You can use recursion to implement this, but the simplest would be just to use itertools.product:

from itertools import product

def f(*sets):
    for p in product(*sets):
        # e.g. (1, 3, 5)
        print(*p)

>>> f([1, 2], [3, 4], [5, 6])
1 3 5
1 3 6
1 4 5
1 4 6
2 3 5
2 3 6
2 4 5
2 4 6

product returns the cartesian product of the input iterables in the form of a lazy iterator over tuples.

A simple recursive implementation of the cartesian product would go along the following lines:

def f(*sets):
    if not sets:
        return [[]]
    result = []
    for head in sets[0]:
        for combo in f(*sets[1:]):
            result.append([head] + combo)
    return result

>>> f([1, 2], [3, 4])
[[1, 3], [1, 4], [2, 3], [2, 4]]

edited Feb 14 '20 at 09:47

answered Feb 14 '20 at 09:30

user2390182

72,016
6
67
89

Thank you very much. I have no idea how to do this with recursion tho, but this will almost work. I didn't notice, but I failed to emphasize how important would this output format be. Would there be a similar generalization if I want to print the products with a comma between them? Such as : "i1,k1,j1" instead of "i1 k1 j1". I'm using this to automatically generate a pseudo AMPL data transforming function, and that syntax is very picky. – Daniel V. Feb 14 '20 at 09:44
I added a recursive approach. What you do inside the innermost for-loop is up to you. You have tuples `p`: you don't need to do `print(*p)`, you can also iterate over them, transform them, or do anything you like. For comma-separated output, you can do `print(','.join(p))`, for instance. – user2390182 Feb 14 '20 at 09:50
Thank you very much for the help. Now I see it way clearer. The recursive snippet will take a while to understand, but that's how my brain works :P Thanks! – Daniel V. Feb 14 '20 at 09:52
Recursion takes a bit to wrap the head around, but once it clicks, you'll love it ;) – user2390182 Feb 14 '20 at 09:52
Using recursion to create the product seems like a bad idea given that it takes about 10 lines of code to solve this in an interative way... – l4mpi Feb 14 '20 at 10:19

Function with varying number of for loops using the arguments in the loop

1 Answers1