python - List comprehension: Returning two (or more) items for each item

100
votes

Is it possible to return 2 (or more) items for each item in a list comprehension?

What I want (example):

[f(x), g(x) for x in range(n)]

should return [f(0), g(0), f(1), g(1), ..., f(n-1), g(n-1)]

So, something to replace this block of code:

result = list()
for x in range(n):
    result.add(f(x))
    result.add(g(x))
7
pythonlist-comprehension
Out of curiosity, why do you want to do this? There might be a better way of achieving your end goal without trying to do it this way. - murgatroid99
Mainly because I like functional programming. I want to map a list of coordinates to a tuple of screen coordinates to use with pyglet.graphics.draw function. - Hashmush
Screen coordinates? As in XY-coordinates? Why do you want to split them up? - blorgon

7 Answers

57
votes
>>> from itertools import chain
>>> f = lambda x: x + 2
>>> g = lambda x: x ** 2
>>> list(chain.from_iterable((f(x), g(x)) for x in range(3)))
[2, 0, 3, 1, 4, 4]

Timings:

from timeit import timeit

f = lambda x: x + 2
g = lambda x: x ** 2

def fg(x):
    yield f(x)
    yield g(x)

print timeit(stmt='list(chain.from_iterable((f(x), g(x)) for x in range(3)))',
             setup='gc.enable(); from itertools import chain; f = lambda x: x + 2; g = lambda x: x ** 2')

print timeit(stmt='list(chain.from_iterable(fg(x) for x in range(3)))',
             setup='gc.enable(); from itertools import chain; from __main__ import fg; f = lambda x: x + 2; g = lambda x: x ** 2')

print timeit(stmt='[func(x) for x in range(3) for func in (f, g)]',
             setup='gc.enable(); f = lambda x: x + 2; g = lambda x: x ** 2')


print timeit(stmt='list(chain.from_iterable((f(x), g(x)) for x in xrange(10**6)))',
             setup='gc.enable(); from itertools import chain; f = lambda x: x + 2; g = lambda x: x ** 2',
             number=20)

print timeit(stmt='list(chain.from_iterable(fg(x) for x in xrange(10**6)))',
             setup='gc.enable(); from itertools import chain; from __main__ import fg; f = lambda x: x + 2; g = lambda x: x ** 2',
             number=20)

print timeit(stmt='[func(x) for x in xrange(10**6) for func in (f, g)]',
             setup='gc.enable(); f = lambda x: x + 2; g = lambda x: x ** 2',
             number=20)

2.69210777094

3.13900787874

1.62461071932

25.5944058287

29.2623711793

25.7211849286

133
votes

Double list comprehension:

[f(x) for x in range(5) for f in (f1,f2)]

Demo:

>>> f1 = lambda x: x
>>> f2 = lambda x: 10*x

>>> [f(x) for x in range(5) for f in (f1,f2)]
[0, 0, 1, 10, 2, 20, 3, 30, 4, 40]
13
votes
sum( ([f(x),g(x)] for x in range(n)), [] )

This is equivalent to [f(1),g(1)] + [f(2),g(2)] + [f(3),g(3)] + ...

You can also think of it as:

def flatten(list):
    ...

flatten( [f(x),g(x)] for x in ... )

note: The right way is to use itertools.chain.from_iterable or the double list comprehension. (It does not require recreating the list on every +, thus has O(N) performance rather than O(N^2) performance.) I'll still use sum(..., []) when I want a quick one-liner or I'm in a hurry, or when the number of terms being combined is bounded (e.g. <= 10). That is why I still mention it here, with this caveat. You can also use tuples: ((f(x),g(x)) for ...), () (or per khachik's comment, having a generator fg(x) which yields a two-tuple).

2
votes

This lambda function zips two lists into a single one:

zipped = lambda L1, L2: [L[i] 
                         for i in range(min(len(L1), len(L2))) 
                         for L in (L1, L2)]

Example:

>>> f = [x for x in range(5)]
>>> g = [x*10 for x in range(5)]
>>> zipped(f, g)
[0, 0, 1, 10, 2, 20, 3, 30, 4, 40]
2
votes

I know OP is looking for a list comprehension solution, but I'd like to offer an alternative using list.extend().

f = lambda x: x
g = lambda x: 10*x

result = []
extend = result.extend
for x in range(5):
    extend((f(x),g(x)))

which is marginally faster than using double list comprehension.

nums = range(100000)

def double_comprehension():
    return [func(x) for x in nums for func in (f,g)]

def list_extend():
    result = []
    extend = result.extend
    for x in nums:
        extend((f(x),g(x)))
    return result

%timeit -n100 double_comprehension()
23.4 ms ± 67 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit -n100 list_extend()
20.5 ms ± 213 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Python version: 3.8.0

0
votes

A solution using reduce:

from functools import reduce

f    = lambda x: f"f({x})" ## Just for example
g    = lambda x: f"g({x})"
data = [1, 2, 3]

reduce(lambda acc, x: acc + [f(x), g(x)], data, [])
# => ['f(1)', 'g(1)', 'f(2)', 'g(2)', 'f(3)', 'g(3)']

This is a functional way of approaching the problem. A list comprehension is essentially another way of maping over data, but in cases such as this one where the mapping isn't one to one between the input and the output, reduce allows some wiggle room with how the output can be generated.

In general, any for implementation of the form:

result = []
for n in some_data:
  result += some_operation()
  ## etc.

(I.e. for loops intended to produce a side effect on a list or similar data structure, the accumulator pattern)

Can be refactored into a declarative map/reduce/filter implementation.

-2
votes

Goodness, grief! Why all these lambdas, flattens, zips and sums? Isn't this the simplest and most readable:

>>> [v
...  for x in range(5)
...      for v in (2 * x,
...                2 * x + 1)]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>>>

(Replace the last two expressions with f(x) and g(x), or whatever.)