Python List Comprehensions – Learn to Write Elegant Python Code

In this tutorial, and in the accompanying YouTube video, we explain how to use Python list comprehensions. List comprehensions enable you to significantly decrease the length of the Python code and to elegantly code. The YouTube video accompanying this post can be found here

Let us assume that you are given the following task:

Write a Python code that will create a list whose every element is a square (second power) of the members of another list.

For example, you are given the following list:

inputList=[1,2,3,…,11]

The output of the code should be

listSquaredNumbers = [1, 4, 9, …, 121]

A non-Pythonic way of doing this (C-style code) is

import numpy as np

listSquaredNumbers=[]

for i in range(12):
    listSquaredNumbers.append(i**2)    


print(listSquaredNumbers)

Another possibility is to define a function that will take powers, and then call the Python map() function. This is illustrated by the following code lines

# another idea without using for loops - use map()

numbers =np.arange(12)

def takeSquares(x):
    return x**2

squaredNumbers=map(takeSquares,numbers)
squaredNumbersList=list(squaredNumbers)
print(squaredNumbersList)

Python’s map() function calls the function takeSequares() for every entry of the list numbers.

However, both of these solutions are relatively cumbersome. Python enables us to significantly compress the code by using list comprehensions. The compact solution takes the following form

# use list comprehensions 
#[expression for member in iterable]

listSquared=[x**2 for x in range(12)]
print(listSquared)

The syntax for the list comprehensions looks like this

“[expression for member in iterable]”

In our case the expression is “x**2”, member is “x” and iterable is “range(12)”.

Another way is to directly call the function “takeSquares()x” from the list comprehensions. Here is the solution

# another way 
listSquared2=[takeSquares(x) for x in numbers]
print(listSquared2)

Here is a code snippet that illustrates the enormous power of list comprehension.

# add if statements 
# [expression for member in iterable (if conditional)]

sentence="My name is John Doe, I am here to help you. Please ask me a question."
vowels='aeiou'

vowelList=[x for x in sentence if x in vowels]

So this code tracks the number of times a vowel appears in the sentence (vowels are a,e,i,o, or u). Here sentence is treated as an array of letters (string). Here we have included an additional functionality of list comprehensions. Namely, we have included if statements. The list will be updated only if “x” is a member of the “vowels” string.

Here are two other nice examples of list comprehensions.

# [expression (if conditional) for member in iterable]

listOfNumbers=[5, -1.5, 7, 0, 2, -1]

listEliminated= [x if x>0 else 0 for x in listOfNumbers]

listEliminated2=[x for x in listOfNumbers if x>0]

Here new lists are created out of the original list “listEliminated” on the basis of the sign of entries of “listEliminated”.

That is all folks for today, make sure to check other Python tutorials I made.