List comprehension and function returning multiple values

Question

I wanted to use list comprehension to avoid writing a for loop appending to some lists. But can it work with a function that returns multiple values? I expected this (simplified example) code to work...

def calc(i):
    a = i * 2
    b = i ** 2
    return a, b

steps = [1,2,3,4,5]

ay, be = [calc(s) for s in steps]

... but it doesn't :(

The for-loop appending to each list works:

def calc(i):
    a = i * 2
    b = i ** 2
    return a, b

steps = [1,2,3,4,5]

ay, be = [],[]

for s in steps:
    a, b = calc(s)
    ay.append(a)
    be.append(b)

Is there a better way or do I just stick with this?

Answer 1

Use zip with *:

>>> ay, by = zip(*(calc(x) for x in steps))
>>> ay
(2, 4, 6, 8, 10)
>>> by
(1, 4, 9, 16, 25)

Answer 2

The horrendous "space efficient" version that returns iterators:

from itertools import tee

ay, by = [(r[i] for r in results) for i, results in enumerate(tee(map(calc, steps), 2))]

But basically just use zip because most of the time it's not worth the ugly.

---

Explanation:

zip(*(calc(x) for x in steps))

will do (calc(x) for x in steps) to get an iterator of [(2, 1), (4, 4), (6, 9), (8, 16), (10, 25)].

When you unpack, you do the equivalent of

zip((2, 1), (4, 4), (6, 9), (8, 16), (10, 25))

so all of the items are stored in memory at once. Proof:

def return_args(*args):
    return args

return_args(*(calc(x) for x in steps))
#>>> ((2, 1), (4, 4), (6, 9), (8, 16), (10, 25))

Hence all items are in memory at once.

---

So how does mine work?

map(calc, steps) is the same as (calc(x) for x in steps) (Python 3). This is an iterator. On Python 2, use imap or (calc(x) for x in steps).

tee(..., 2) gets two iterators that store the difference in iteration. If you iterate in lockstep the tee will take O(1) memory. If you do not, the tee can take up to O(n). So now we have a usage that lets us have O(1) memory up to this point.

enumerate obviously will keep this in constant memory.

(r[i] for r in results) returns an iterator that takes the i^th item from each of the results. This means it receives, in this case, a pair (so r=(2,1), r=(4,4), etc. in turn). It returns the specific iterator.

Hence if you iterate ay and by in lockstep constant memory will be used. The memory usage is proportional to the distance between the iterators. This is useful in many cases (imagine diffing a file or suchwhat) but as I said most of the time it's not worth the ugly. There's an extra constant-factor overhead, too.

Answer 3

You should have shown us what

[calc(s) for s in xrange(5)]

does give you, i.e.

[(0, 0), (2, 1), (4, 4), (6, 9), (8, 16)]

While it isn't the 2 lists that you want, it is still a list of lists. Further more, doesn't that look just like?

zip((0, 2, 4, 6, 8), (0, 1, 4, 9, 16))

zip repackages a set of lists. Usually it is illustrated with 2 longer lists, but it works just as well many short lists.

The third step is to remember that fn(*[arg1,arg2, ...]) = fn(arg1,arg2, ...), that is, the * unpacks a list.

Put it all together to get hcwhsa's answer.