CODEWITHSUNDEEP
pythonpython-3.x
Daniel V.
Daniel V.

Reputation: 87

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

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

Upvotes: 2

Views: 396

Answers (1)

user2390182
user2390182

Reputation: 73460

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]]

Upvotes: 5

Related Questions