'List comprehension: Returning two (or more) items for each item
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))
Solution 1:[1]
>>> 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
Solution 2:[2]
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]
Solution 3:[3]
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).
Solution 4:[4]
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
Solution 5:[5]
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]
Solution 6:[6]
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 map
ing 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.
Solution 7:[7]
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))]
Solution 8:[8]
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.)
Sources
This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.
Source: Stack Overflow
Solution | Source |
---|---|
Solution 1 | |
Solution 2 | |
Solution 3 | |
Solution 4 | remykarem |
Solution 5 | Daniel Reis |
Solution 6 | |
Solution 7 | |
Solution 8 | Guest |