Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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.

like image 316
Christophe Avatar asked Feb 21 '17 19:02

Christophe


1 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).

like image 196
Willem Van Onsem Avatar answered Oct 24 '22 19:10

Willem Van Onsem