Search code examples
pythonpython-3.xlist-comprehension

How can I get multiple lists as separate results from a list comprehension?


Suppose I have this code:

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

x = range(3)
xlist, ylist = [f(value) for value in x]

How can I neatly get a result like this?

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

Solution

  • Note that return 2*x,x is short for return (2*x,x), i.e. 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]
    

    Another way to do this is with separate list comprehensions:

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

    Of course, this repeats the work of f, which is inelegant and can be inefficient.