George V. Reilly

Flattening List Comprehensions in Python

List Comprehension

Python has list com­pre­hen­sions, syntactic sugar for building lists from an expression.

>>> [2 * i for i in (2, 3, 5, 7, 11)]
[4, 6, 10, 14, 22]

This doesn’t work so well when the com­pre­hen­sion expression is itself a list: you end up with a list of lists.

>>> def gen():
...     for l in [['a', 'b'], ['c'], ['d', 'e', 'f']]:
...         yield l
...
>>> [l for l in gen()]
[['a', 'b'], ['c'], ['d', 'e', 'f']]

This is ugly. Here’s one way to build a flattened list, but it’s less elegant than the com­pre­hen­sion.

>>> x = []
>>> for l in gen():
...     x.extend(l)
...
>>> x
['a', 'b', 'c', 'd', 'e', 'f']

It took me a while to find a readable list com­pre­hen­sion, with a little help from Google. Use sum() on the outer list and prime it with an empty list, []. Python will con­cate­nate the inner lists, producing a flattened list.

>>> sum([l for l in gen()], [])
['a', 'b', 'c', 'd', 'e', 'f']

Al­ter­na­tive­ly, you can use itertools.chain().

>>> import itertools
>>> list(itertools.chain(*gen()))
['a', 'b', 'c', 'd', 'e', 'f']

That might be slightly more efficient, though I find the sum() to be a little more readable.

>>> import itertools
>>> list(itertools.chain(*gen()))
['a', 'b', 'c', 'd', 'e', 'f']

That might be slightly more efficient, though I find the sum() to be a little more readable.

Edit: I forgot about nested com­pre­hen­sions

>>> [inner
...     for outer in gen()
...         for inner in outer]
['a', 'b', 'c', 'd', 'e', 'f']

Somewhat cryptic on one line however:

>>> [j for i in gen() for j in i]
['a', 'b', 'c', 'd', 'e', 'f']

Update: The nested com­pre­hen­sion became one of my most popular Stack­Over­flow answers.

blog comments powered by Disqus
Circe, part 1 » « Review: The Choirboys