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

Viewed 45850

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))
8 Answers

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

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.

When reasonably possible, you should prefer using itertools.product() over using two for clauses in a list comprehension. Some style guides, such as the Google Python Style Guide forbid using multiple for clauses in a list comprehension. While you may be using a style guide that doesn't explicitly forbid this practice, you can still make the comprehension more concise by doing something like the following:

from itertools import product

result = [y(x) for y, x in product(range(n), (f, g))]

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

Related