Advanced listcomps

A week ago, @alex_gaynor tweeted a quick listcomp for flattening a two-dimensional matrix (list of lists):

At the time, I thought it was cool, but I didn’t really “get” the “how”. But now I do.

Say we have the following matrix:

matrix = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
]

The flattening looks like this:

>>> [x for y in matrix for x in y]
[1, 2, 3, 4, 5, 6, 7, 8, 9]

PEP 202 states:

  • The form [… for x… for y…] nests, with the last index varying fastest, just like nested for loops.

Most list comps take the form a = [b(c) for c in d if e]. If each level of logic nests, this is equivalent to:

a = []
for c in d:
    if e:
        a.append(b(c))

So, following this logic, multiple for ... clauses in the matrix-flattening are simply nested loops:

matrix = [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
]

flattened = [x for y in matrix for x in y]

flattened = []
for y in matrix:
    for x in y:
        flattened.append(x)