List comprehensions#
I use list comprehensions extensively throughoput the book. Its essential that you understand how these work. List comprehensions are a compact alternative to for loops. They work specifically with Python Lists.
Example 1 - double the numbers#
If we take the following for loop that doubles the numbers in foo
foo = [1, 2, 3, 4]
# `bar` is the list we will fill.
bar = []
# loop through all values in foo and append to bar
for x in foo:
bar.append(x * 2)
print(bar)
[2, 4, 6, 8]
The list comprehension that produces the same results is
foo = [1, 2, 3, 4]
bar = [x * 2 for x in foo]
print(bar)
[2, 4, 6, 8]
Looking at this in more detail we can see that as in the first foo is a pre-populated list.
Instead of of a standard for loop the variablebar is assigned a value using a simple list notation. This is essentially a list with a for loop inside of it. The generic notation is:
my_list = [<operation on x> for x in foo]
We don’t need to call the current item x. That is just an arbitray choice for the name of the variable. We could for example use i, a or whatever a more desriptive name if you feel is appropriate.
Note how compact the list comprehension code is relative to the for loop. Its less code and we didn’t need to declare bar as an empty list in advance.
The other thing to say is that foo doesn’t need to be a list. We could use any iterable object. For example a range. We could modify the original example to:
# `bar` is the list we will fill.
bar = []
# loop through all values in the range 1 to 4
for x in range(1, 5):
bar.append(x * 2)
print(bar)
[2, 4, 6, 8]
The equivalent list comprehension is:
bar = [x * 2 for x in range(1, 5)]
print(bar)
[2, 4, 6, 8]
I find list comprehensions like the above very readable and good for mathematics.
Example 2 - Using a function inside a comprehension#
The operation we used above was specified inside list comprehension. If we had a more complex task, for example converting celsius to fahrenheit we could call a function from within the comprehension. Let’s try that out. Here is our function:
def convert_celsius_to_fahrenheit(deg_celsius):
"""
Convert degress celsius to fahrenheit
Returns float value - temp in fahrenheit
Parameters:
-----------
def_celcius -- temp in degrees celsius
"""
return (9/5) * deg_celsius + 32
# list of temps in degree celsius to convert to fahrenheit
celsius = [39.2, 36.5, 37.3, 41.0]
We call the function list so:
converted = convert_celsius_to_fahrenheit(celsius[0])
print(converted)
102.56
If we wanted to convert all our data from celsius to fahrenheit we could use a standard for loop as following:
fahrenheit = []
for x in celsius:
fahrenheit.append(convert_celsius_to_fahrenheit(x))
print(f'using standard for loop: {fahrenheit}')
using standard for loop: [102.56, 97.7, 99.14, 105.8]
But equally we can create more compact code using list comprehension syntax
fahrenheit = [convert_celsius_to_fahrenheit(x) for x in celsius]
print(f'using list comprehension: {fahrenheit}')
using list comprehension: [102.56, 97.7, 99.14, 105.8]
This is effectively the same pattern as in our double the numbers example. We just replaced x * 2 with convert_celsius_to_fahrenheit(x)
Example 3 - Using if statements within a list comprehension#
You can extend list comprehensions to include if statement. This limits what is used in the loop. I often use this as a simple approach to filter a small list. A common bit of code I use is filtering file types.
This example filters a list of file names to the python files only. First here is the standard for loop example. Note the inclusion of the if conditional to identify the python files.
unfiltered_files = ['test.py', 'names.csv', 'fun_module.py', 'prog.config']
python_files = []
# filter the files using a standard for loop
for file in unfiltered_files:
if file[-2:] == 'py':
python_files.append(file)
print('using standard for loop: {}'.format(python_files))
using standard for loop: ['test.py', 'fun_module.py']
The list comprehension code if far more compact. We simpy add the if statement onto the end of the syntax. I.e.
python_files = [file for file in unfiltered_files if file[-2:] == 'py']
print(f'using list comprehension {python_files}')
using list comprehension ['test.py', 'fun_module.py']
Example 4 - List comprehension to create a list of lists#
Its likely at some point that you will need a list of lists. List comprehension syntax here is very flexible and will allow you to do create these lists, but its a bit harder to follow in my opinion. This is equivalent to a nested for loop. Here is some standard code to create one:
list_of_lists = []
# outer loop: controls how many items are in the outer list
for i in range(5):
# on each iter we create a new sublist to nest
sub_list = []
# inner loop: controls how many items are in the sub list.
for j in range(3):
# arbitrary operation for the example
sub_list.append(i * j)
# we append the sunlist to the list of lists.
list_of_lists.append(sub_list)
print(list_of_lists)
[[0, 0, 0], [0, 1, 2], [0, 2, 4], [0, 3, 6], [0, 4, 8]]
a list comprehension reduces 6 lines of code to 1!
list_of_lists = [[i * j for j in range(3)] for i in range(5)]
print(list_of_lists)
[[0, 0, 0], [0, 1, 2], [0, 2, 4], [0, 3, 6], [0, 4, 8]]
Although the reduction in code is nice. It is less readable at first glance. Let’s break it down.
The first thing to note is that we have an inner and outer list comprehension just list the nested for loops. Let’s simplify to see it:
[[<inner comprehension>] for i in range(5)]
When we look at the code list this it doesn’t seem so bad does it? Its just the same as examples 1 and 2. But we replace the mathematical operation or function call with another list comprehension. Just like function calls we can use i (our current variable in the loop) in the the inner list comprehension. i.e.
[i * j for j in range(3)]
Here j varies, but i remains constant. When we put these two together we get the list of lists.
Example 5: Iterate over all items in a list of lists#
This time we will start with a list of lists. How do we iterate over all items in the list?
As an example use this list
[[8, 2, 1], [9, 1, 2], [4, 5, 100]]
list_of_lists = [[8, 2, 1], [9, 1, 2], [4, 5, 100]]
flat_list = []
for sublist in list_of_lists:
for item in sublist:
flat_list.append(item)
print(flat_list)
[8, 2, 1, 9, 1, 2, 4, 5, 100]
The code above iterates through each item in turn e.g. 8, 2, 1, 9 … 100. At the end of the snippet we have produces a flat list (1 dimension). But we need not have created a list we could have called a function or performed some processing inline.
To create a flat list with list comprehension syntax we can use the following code:
list_of_lists = [[8, 2, 1], [9, 1, 2], [4, 5, 100]]
flat_list = [item for sublist in list_of_lists for item in sublist]
print(flat_list)
[8, 2, 1, 9, 1, 2, 4, 5, 100]
I’ve used the same variable names as the for loop code. I.e. sublist and item. I think that makes the comparison of the code easier.
You can think of the nested comprehension in two parts. This is the same as the standard for loop i.e. outerloop then inner loop.