list comprehension for multiple return function?

I have a function that returns two values, and I would like use a list comprehension to fill two lists. for example:

def f(x): return 2*x,x*x x = range(3) xlist, ylist = [f(value) for value in x] EDITS from answers below: xtuple, ytuple = zip(*[f(value) for value in x]) xlist, ylist = map(list,zip(*[f(value) for value in x])) 

where the expected return should be:

xlist = [0, 2, 4] ylist = [0, 1, 4] 

my question boils down to:

Currently I get a list of tuples, while this is reasonable I will end up needing two independent lists. currently I could have 1 placeholder (tuple list) variable, and 3 total comprehensions. But I'm wondering if there is a clean way to do it with as single list comprehension.

Worth noting: in the real code my two returns are correlated, so I cannot simply split the function into two.

9

6 Answers

First of all, you made a small mistake: it should be:

[f(value) for value in x] # ^ notice the `value` 

instead of:

[f(x) for value in x] 

Furthermore the point is that:

return 2*x,x 

is short for:

return (2*x,x) 

so a tuple. Your list comprehension thus generates a list of tuples, not a tuple of lists. The nice thing of zip however is that you can easily use it in reverse with the asterisk:

xlist,ylist = zip(*[f(value) for value in x]) # ^ with asterisk

Note that xlist and ylist will be tuples (since zip will be unpacked). If you want them to be lists, you can for instance use:

xlist,ylist = map(list,zip(*[f(value) for value in x]))

which results in:

>>> xlist [0, 2, 4] >>> ylist [0, 1, 4] 

(note that ranges start counting from 0)

Alternative: Another way to do this is of course:

xlist = [f(value)[0] for value in x] ylist = [f(value)[1] for value in x] 

But this is of course inelegantly and furthermore can be inefficient (given f is computationally expensive).

2

Let's make this work. The function is fine:

def f(x): return 2*x, x*x 

But you want to define the range as follows, notice the starting and ending values:

x = range(1, 4) 

Also, you have to call the function with the value, not with the list as parameter. And the final trick to unzip the result into two lists, is to simply zip(*lst) the result of the list comprehension:

xlist, ylist = zip(*[f(value) for value in x]) 

Now the result is as expected:

xlist => [2, 4, 6] ylist => [1, 4, 9] 

Use the build-in function zip(),

def f(x): return 2*x, x*x x = range(1, 4) xlist, ylist = zip(*[f(value) for value in x]) print(xlist, ylist) # ((2, 4, 6), (1, 4, 9)) 

Use

zip(*your_list_of_bituples) 

Example

demo_list = [(1, 2), (2, 3), (4, 5)] zip(*demo_list) 

Will give

[(1, 2, 4), (2, 3, 5)] 
1

I know it's late but the following gets what you want.

def f(value): xlist = [] ylist = [] for x, y in [(2*x, x*x) for x in range(value)]: xlist.append(x) ylist.append(y) return xlist, ylist x = int(input("enter a value: ")) xval, yval = f(x) print(f"xlist = {xval}\nylist = {yval}") 
def f(x): yield [2*x, x*x] xlist, ylist = zip(*[next(f(x)) for x in range(3)]) print(list(xlist)) print(list(ylist)) 

using yield...

[0, 2, 4] [0, 1, 4] [Program finished] 

Your Answer

Sign up or log in

Sign up using Google Sign up using Facebook Sign up using Email and Password

Post as a guest

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

You Might Also Like