Introduction
If you are a python programmer you probably love list comprehensions and you use them a lot in your code. However, I have noticed that many beginner Data Scientists do not know what list comprehensions are or are confused about how to use them properly.
This article aims to explain what are list comprehensions and why you should use be using them.
It will also give you practical examples where list comprehension can be used.
What are list comprehensions and why you should use them?
List comprehensions are code one-liners that allow you to create lists from Iterables (for example lists).
Do you remember for loops that you have been creating to iterate through lists? In some cases, they can be replaced by list comprehensions, and let’s see an example of it. Imagine that you have a list of numbers and you want to multiply each element in the list by number 2. Below you will see two different ways of doing this operation.
nums = [1, 2, 3, 4]
The first piece of code is written using traditional for loop and the second one is written using list comprehension.
So which one do you prefer a traditional loop, or list comprehension version?
I think it is clear that the second piece of code is cleaner, looks more elegant, and is more ‘pythonic’ as some people would say. Those should be enough reasons to use list comprehensions rather than traditional for loops.
Additionally, list comprehensions will usually execute faster and use less memory!
List comprehension syntax
Let’s look at a basic list comprehension syntax a bit closer.
First, note that the whole expression is in square brackets [ ].
After bracket opening ( [ ), we have an expression that is applied to every element of the list: num * 2.
This is followed by the loop written on the same line: for num in nums.
There is a bracket closing to finish the whole expression ( ] ).
And that is all!
Simple, no?
This why the list comprehensions are so popular with programmers, elegant, simple, and concise.
List comprehension with a condition
You have learned basic list comprehension syntax but there is some additional syntax that could be used to filter a list that is being created. You can add if statement just after the for loop to include a condition that must be met for new list elements.
Let’s look at some examples:
one condition
nums = [1, 2, 3, 4]
doubled_nums = [num * 2 for num in nums if num % 2 == 0]
[4, 8]
Here we have our previous example but it is slightly modified. We have added a condition after the for loop that states that the numbers that we will be taking from the initial list must be divisible by 2.
This is achieved by adding the following condition: num % 2 == 0.
As you can see the result has only two elements in the list as we only retained even numbers from the initial list.
double condition
You can have more than one condition. In this case, multiple conditions can be combined using ‘and’ operator.
nums = [1, 2, 3, 4]
doubled_nums = [num * 2 for num in nums if num % 2 == 0 and num > 3]
[8]
Here we are filtering our list to only include numbers that are divisible by 2 and are greater than 3.
The result of the filtering is only one number: 4. That number is then multiplied by 2 and we end up with one element list: [8].
List comprehension with a nested loop
Now that you have learned basic syntax and you know how to filter the list, we are going to learn how to use an additional loop. This will allow us to work with nested lists.
Yes, list comprehensions can handle nested lists by using nested loops!!!
Imagine that you have a list of lists:
nested_list = [[1,2], [3,4], [4,5]]
Now you want to take all elements out of it and put in a single list. This is called flattening a list.
flattened_list = [item for sublist in nested_list for item in sublist]
[1, 2, 3, 4, 4, 5]
Can you see how we nested the loops? The order is the same as with a traditional loop. Below there is the same code written with traditional for loop syntax:
flattened_list = []
for sublist in nested_list:
for item in sublist:
flattened_list.append(item)
[1, 2, 3, 4, 4, 5]
List comprehensions with other Iterables
So far you have been using list comprehensions based on some initial list. The great news is that the same as with traditional loops the underlying data does not need to be a list. It can actually be any Iterable python object.
This is what makes list comprehensions so powerful. Let’s look at examples:
list comprehensions with strings
string = 'boat'
letters = [letter for letter in string]
letters
['b', 'o', 'a', 't']
The code above lists all the letters in a string ‘boat’.
list comprehensions with dictionaries
numbers_dictionary = {1: 'one', 2: 'two', 3: 'three'}
multiplied_keys = [k * 2 for (k,v) in numbers_dictionary.items()]
multiplied_keys
[2, 4, 6]The code above creates a list that consists of all dictionary keys multiplied by a factor of 2.
When not to use list comprehensions?
We have talked a lot about the advantages of list comprehensions but is there any time you should not use them?
The answer is: Yes.
If the code is too complex sticking everything in one-liner is not the best option, even if it works.
You should go for readability and if there are too many operations to be done the list comprehension will look messy and be really hard to read. In this case, you should use traditional loops.
Summary and call to action
So now you know the power of list comprehensions. It is time to take action and challenge yourself.
From now on, every time you are about to use and old fashioned for loop, STOP, and use list comprehension instead!