Python List Comprehension Part Two

In last weeks post about list comprehension I gave you an overview and some simple examples of how this Pythonic convention works.

There’s so much more you can do with list comprehensions – today I’ll try and give you some examples that are a little more involved.

More Advanced List Comprehension

Last week I wrote about replacing for loops with list comprehensions. If you are an experienced coder with limited exposure to this Python paradigm, you may well have asked yourself “what about nested loops?”.

While not every for loop can be replaced with a comprehension, nested loops are not a problem.  Lets say that we have a matrix ( or a list of lists, or whatever it is that you call this concept) like this:

mat = [[1,2,3,4], [5,6,7,8], [9,10,11,12], [13,14,15,16]]

If we needed to flatten this matrix while keeping only the even numbers.  You could certainly do this the old school way and nest a loop in another loop:

# For Loop Version

result = []

for row in mat:
    for i in row:
        if i%2 == 0:


# [2, 4, 6, 8, 10, 12, 14, 16]

This works fine, no problem.  However it is a bit verbose and not Pythonic at all.  Here is a more economical, Pythonic, way to accomplish the same task:

# List Comprehension version

[i for row in mat for i in row if i%2 == 0]

# [2, 4, 6, 8, 10, 12, 14, 16]

Cool, right?  Don’t get overwhelmed, if this seems complicated.  If you read my first post about list comprehension you already know how this works.

if i%2==0

This is simply the the third ( optional ) part of a list comprehension the Conditional Filtering.  Here we filter so we are only keeping only the numbers that are mod 2 == 0 or even.

The magic in this comprehension is the:

i for row in mat for i in row

You see how we are saying for row in mat for i in row?  Those are not 2 separate loops – they are how we nest loops in a list compression!

But be careful.  This can get wacky quickly if small errors are introduced:

[i for row in row if i%2==0]
# [16, 16, 16, 16]

What if, instead we had two separate ( not nested ) lists:

list_a = [9, 3, 6, 1, 5, 0, 8, 2, 4, 7]
list_b = [6, 4, 6, 1, 2, 2]

And we wanted to return a list of tuples where each tuple contained a number from list_b and it’s corresponding position in list_a:

# Old school loop solution

result = []

for i in list_b:
    result.append((i, list_a.index(i)))

# [(6, 2), (4, 8), (6, 2), (1, 3), (2, 7), (2, 7)]

# the same as above with list comp
[ (i, list_a.index(i)) for i in list_b ]

# [(6, 2), (4, 8), (6, 2), (1, 3), (2, 7), (2, 7)]

See those parenthesis around our 2 returned values?  We need them for tuples of course – and it illustrates how you can add other characters in your comprehensions.

You can do the same kind of thing with dictionaries:

{i: list_a.index(i) for i in list_b}
# {6: 2, 4: 8, 1: 3, 2: 7}


What if we wanted to make a new list containing all of the numbers in both lists, but we wanted each number to occur only once?  No problem, just grab them all in a set and cast to a list:

list({a for a in list_a for b in list_b})
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

# and then only numbers in both lists

list({a for a in list_a for b in list_b if a == b})
# [1, 2, 4, 6]

More List Comprehensions

My research has suggested that the optimal number of words for an educational post is 500.  I have reached that count at this point – but there is so much more you can do with comprehensions!

You can see all the code I created for study on list comprehension in a Jupyter Notebook here

The notebook contains more code than I’ve included in these first two posts.  I plan on writing a part three list comprehension post – but if I don’t lol, download that notebook and get the rest of my examples.

1 comment

Leave a comment

Your email address will not be published. Required fields are marked *