Python – List Comprehensions

Posted on 09-17-2018 by superadm1n

Card image cap
In Python you can use list comprehensions to shorten the amount of code that you need to use to accomplish tasks when working with lists. The first time I heard of list comprehensions it sounded very advanced and confusing and when I looked at some example code I was, as I expected confused. Fear not list comprehensions are just as simple as a for loop and once you understand them you will begin to use them in your code to solve problems in much more elegant ways and conserving unneeded lines of code.

 

What is a List Comprehension and What is its Purpose?

A list comprehension is a concept in Python that’s used to generate a list. Its purpose is to allow a programmer to generate a list in a more fluent and efficient way and use less code in the process than if you were to use a for loop to generate a list. If you do not understand what a List is in python check out my post linked here to become more familiar with lists.

List Comprehension Syntax:

You can think of a list comprehension as a single line for loop enclosed in brackets (thats because it pretty much is) with the operation your performing coming first and the for loop statement coming second. Below I have a visual representation of a traditional for loop and where you place it in a list comprehension

As you can see the code in the blue square (the operation your performing) comes first inside the brackets and the code in the red square (for loop statement minus the colon “:”) comes second.

With this visual representation I hope it is easy to see that a list comprehension can be thought of as a reverse single line for loop.

 

Example Scenario:

Lets say we have a list of numbers and we need to add 5 to each of those numbers and print them out to the user. Below I have an example using a for loop to accomplish that task and then I solve the same problem using a list comprehension.

For Loop:

#!/usr/bin/python3

# List of numbers to be displayed to the user
list_of_numbers = [2, 5, 16, 20]

# Define new list
list_to_be_displayed = []

# For loop to add 5 to each of the numbers in my list
for number in list_of_numbers:
    list_to_be_displayed.append(number + 5)

# Prints the new numbers out to the user
for number in list_to_be_displayed:
    print(number)

Now I will be the first to admit there are more elegant ways to solve this without the need to use 2 for loops, for instance I could have added print(number + 5) inside the first for loop to remove the need for a second for loop but at the same time I would be doing the math twice vs a second for loop. Now lets solve the same problem using a list comprehension.

List Comprehension:

#!/usr/bin/python3

# List of numbers to be displayed to the user
list_of_numbers = [2, 5, 16, 20]

# Adds 5 to each number in the list of numbers
list_of_numbers = [number + 5 for number in list_of_numbers]

# prints out each number to the user
[print(number) for number in list_of_numbers]

Running both of the code snippets gives us the following output:

User@Laptop> forloop.py
7
10
21
25
User@Laptop>

The mechanics of what is happening in both examples is almost the exact same the only exception is we are not creating a new list for the new numbers, instead the list comprehension at line 7 allows us to overwrite our existing list variable with the new values. We have reduced our total overall code by 5 lines but we still use 2 list comprehensions to first do the addition to each of the numbers and the second to print them to the user because we have kept our for loops to a single line.

Generate Strings with List Comprehensions:

Another useful way to use list comprehensions is when generating strings using values in a list, below I have an example based on the original one above to create a human readable string of the added numbers to display to the user that reads “My following numbers are 7 and 10 and 21 and 25”.

Using For Loop:

#!/usr/bin/python3

# List of numbers to be displayed to the user
list_of_numbers = [2, 5, 16, 20]

#creates string variable
string = ''

# takes our list, adds 5 to each number and appends that number with the word " and " to the string
for number in list_of_numbers:
    string += '{} and '.format(number + 5)
# Removes the last word " and " from our string
string = string[:-5]
# Prints out to the user
print('My following numbers are {}'.format(string))

When using a for loop we have to create a new variable for the string, then we have to use a for loop to generate the custom part of our string (## and ## and etc.) then we need to remove the final and from the custom part of our string before we can finally print it to the user

Using List Comprehension:

#!/usr/bin/python3

# List of numbers to be displayed to the user
list_of_numbers = [2, 5, 16, 20]

# prints out each number to the user
print('My following numbers are {}'.format(' and '.join([str(number + 5) for number in list_of_numbers])))

Using a list comprehension we have shaved 8 lines off of our code and done all of the work in a single line. Now I will be the first to admit to understand the mechanics of what is happening here you need to analyze line 7 alittle harder but once you start getting comfortable reading more complex list comprehensions that one line it is easy to see what is happening.

Here is the breakdown of what is happening in line 7:

First the list comprehension is generating a new list and that list will contain the product of each of the numbers and 5 storing it as a string in the list that is being generated. That code is the list comprehension itself.

[str(number + 5) for number in list_of_numbers]

Second that list of strings is put together in the following format “## and ## and ## and ##”. That is done by using the join() method which has the list comprehension nested inside of it. Note the join() method is smart enough not to add the word “and” after the last element in the list so we don’t need to handle it like we did before.

' and '.join([str(number + 5) for number in list_of_numbers])

At this point we have our custom string generated, formatted, and ready to be attached to the static part of our string.

Third our custom string is joined with the static part of our string by using the format() method with this code. (If you would like to understand more on string formatting check out my post linked here)

'My following numbers are {}'.format(' and '.join([str(number + 5) for number in list_of_numbers]))

And finally it is printed out to the user by encasing all of that within a print statement.

print('My following numbers are {}'.format(' and '.join([str(number + 5) for number in list_of_numbers])))

Running both of the code snippets gives us the following output:

User@Laptop> test.py
My following numbers are 7 and 10 and 21 and 25
User@Laptop>

Nested List Comprehensions:

You can also nest list comprehensions much like you nest for loops so for instance if you have a list of lists and you need to run a list comprehension over each element in the nested lists this can be accomplished by using a list comprehension. Below I show you an example of adding 1 to each element in a list of lists, the first block I will show you using a for loop and the second i will show you using a list comprehension.

#!/usr/bin/python3

# generates list and prints it out
mylist = [[1, 2, 3], [4, 5, 6]]
print('Original List: {}'.format(mylist))

# adds 1 to each element of each list within the main list
for list in mylist:
    counter = 0
    for element in list:
        addition = element + 1  # adds 1 to the current element
        list.pop(counter)  # removes current element
        list.insert(counter, addition)  # inserts new number
        counter += 1  # increments counter

# prints list
print('Modified List: {}'.format(mylist))
#!/usr/bin/python3

# generates list and prints it out
mylist = [[1, 2, 3], [4, 5, 6]]
print('Original List: {}'.format(mylist))

# adds 1 to each element of each list within the main list and prints list
mylist = [[element + 1 for element in list] for list in mylist]
print('Modified List: {}'.format(mylist))

Running both of these code examples gives us the following result.

User@Laptop> test.py
Original List: [[1, 2, 3], [4, 5, 6]]
Modified List: [[2, 3, 4], [5, 6, 7]]
User@Laptop>

As you can see running each snippet of code they yield the same result but using a list comprehension takes off 8 lines of code that is the nested for loop and does everything in one line. Using the list comprehension we do not have to create any intermediate variables to count the number of loops or hold temporary data we simply re declare the same variable of our list “mylist” and set it equal to the list comprehension and it takes care of all of the heavy lifting for us which allows us to have cleaner and more understandable code.

If understanding what is happening when translating a nested for loop into a list comprehension is hard for you to see below I have another visual representation as to where each part of the for loop goes into a list comprehension.

Looking at the picture we can see that in the list comprehension we have 2 sets of square brackets with the outer most of the for loop (red) in the outer most bracket, and the inner for loop (blue) inside the inner brackets with the operation we are performing (green) prepended to the beginning inside the inner brackets.

I hope that this article has cleared up any mystery with list comprehensions and I was able to provide some useful examples. If you have any questions please drop a comment below and I will be sure to get back to you with an answer as soon as possible and if this article was helpful please share on social media!

 

 

 

 

 

 

 

 

 


Leave a Comment:

Search