Python List Comprehension Tutorial

What is Python's List Comprehension?

In a nutshell python's list comprehension is just some syntactic sugar to solve an already existing problem. 

What is this problem? The problems this solves is how to transform one iterable, such as a list or dictionary, into another iterable. 

People from a mathematics background will find some familiarity in the design of list comprehension. An example is {X^2: X in [0,1,2,3,4]} which basically gets the set of numbers between 0 and 4 inclusive squared ie. (0,1,4,9,16).

We can do this in python via the following command:

squared = [x**2 for x in range(0,5)] 
# [0, 1, 4, 9, 16]

Conditions in List Comprehension

What if in the above example you only wanted to get a list of numbers that were numbers of leap years and less then or equal to 2000?

List comprehension can handle this!

First we will need to create a function that determines whether a number is a leap year. Note the algorithm for a leap year is: 

if (year is not divisible by 4) then non leap year
else if (year not divisible by 100) then leap year
else if (year not divisible by 400) then non leap year
else leap year

The function we will use is below:

def leap_year(num):
if (num%4 != 0):
return False
elif (num%100!=0):
return True
elif (num%400!=0):
return False
else:
return True

Now lets get the squared numbers of interest via list comprehension

leap_year_squared = [x**2 for x in range(0,2001) if leap_year(x)]
leap_year_squared
# [0,16,...,3984016,4000000]

Lists vs Dictionary

Why would you choose a List over a Dictionary?

The reason for choosing one data structure over another is whether you want data to be repeated or now. Lists allow duplication, Dictionary do not. Lets see this in action.

numbers = [1,1,2,3,4,4,4,5,6,7,7,7,6]
ls = [x for x in numbers]
ls
# [1, 1, 2, 3, 4, 4, 4, 5, 6, 7, 7, 7, 6]
dict = {x for x in numbers}
dict
# {1, 2, 3, 4, 5, 6, 7}

Nested Loops

Lets go through an example of how you can use list comprehension and nested loops. Lets use list comprehension to extract all elements of a list of lists (or a matrix for this case) into a list.

matrix = [[1,2,3,],[4,5,6],[7,8,9]]
sum_list = [element for row in matrix for element in row]
sum_list
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

Conclusion

Hopefully you will see how useful list comprehension is in python. It enables you to write quite complex logic into one liners. Leave a comment of any cool list comprehensions you have come across and would like to share with others.

Subscribe to our mailing list

* indicates required