Flattening List Comprehensions in Python
Python has list comprehensions, 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 comprehension 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 comprehension.
>>> 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 comprehension, with a little help from Google. Use sum() on the outer list and prime it with an empty list, []. Python will concatenate the inner lists, producing a flattened list.
>>> sum([l for l in gen()], []) ['a', 'b', 'c', 'd', 'e', 'f']
Alternatively, 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 comprehensions
>>> [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 comprehension became one of my most popular StackOverflow answers.