calling function in a loop for several lists python

Viewed 49

I have written a function in python which takes three input variables. Now I want to call the function on a list of items for each input variables. So the function call is like this:

 call_function('cat', 'brown', 2)

The lists are:

 category_list = ['cat', 'dog', 'squirrel']
 color_list = ['brown', 'black']
 age = [1,2,3]

Now I want to run function on each item. So it should return results for cat, brown and 1, and then cat, brown and 2, and then cat, brown and 3 and then cat, black and 1, and then cat, black and 2, and then cat, black and 3 and then do the same for dog and then for squirrel. I know that I will need to call the function in a loop, but I tried and failed several times. Could someone please help me with this?

3 Answers

use itertools, it does the job for you, really simple to call the product function.

supply the variables to itertools.product(....)) as paramaters then list() the response.

Code

import itertools
category_list = ['cat', 'dog', 'squirrel']
color_list = ['brown', 'black']
age = [1,2,3]

def call_function(category, color, age):
    print(category, color, age)


combinations = list(itertools.product(category_list, color_list, age))
for combintaion in combinations:
    call_function(combintaion[0], combintaion[1], combintaion[2])

Output

cat brown 1
cat brown 2
cat brown 3
cat black 1
cat black 2
cat black 3
dog brown 1
dog brown 2
dog brown 3
dog black 1
dog black 2
dog black 3
squirrel brown 1
squirrel brown 2
squirrel brown 3
squirrel black 1
squirrel black 2
squirrel black 3

Proof

Test Here enter image description here

code for method from itertools

def product(*args, repeat=1):
    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
    pools = [tuple(pool) for pool in args] * repeat
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)

here is another way to do it, using pd.apply, after creating a dataframe using itertools.product

import itertools  

category_list = ['cat', 'dog', 'squirrel']
color_list = ['brown', 'black']
age = [1,2,3]

df=pd.DataFrame.from_records(
    list(itertools.product(
        category_list, color_list, age
    )), columns=['category', 'color', 'age'])

df.apply(lambda row: call_function(row['category'], row['color'], row['age']), axis=1)

RESULT:

cat brown 1
cat brown 2
cat brown 3
cat black 1
cat black 2
cat black 3
dog brown 1
dog brown 2
dog brown 3
dog black 1
dog black 2
dog black 3
squirrel brown 1
squirrel brown 2
squirrel brown 3
squirrel black 1
squirrel black 2
squirrel black 3

if function is named MY_FUNCTION

df.apply(lambda row: MY_FUNCTION(row['category'], row['color'], row['age']), axis=1)

enter image description here

I solved it by calling the function in a loop. I am sure there is a more efficient/elegant way to do this, but since my df and lists are small, it works good. Wanted to post my solution for those who are working with small lists and having similar problem. Of course, improvement suggestions are welcome.

for i in category_list: 
    for j in color_list: 
        for k in age: 
            call_function(i,j,k)
Related