5.4. List comprehensions

The loops presented in the previous part of this chapter are fundamental operations that are common in most programming languages. Their syntax and usage at this point is hopefully somewhat simple to understand, and you can get by perfectly well using them in 99% of situations. However, in Python, there is a somewhat odd way of wrapping up the for loop structure into a list as a method for easily creating a list that contains the results of simple calculations performed in an iteration.

Consider the following situation: you want to generate a list containing the cube root of each number in a list. We could do this using a for loop as below:

mylist = [1, 2, 3, 4, 5, 6, 7]
cube_roots = []

for element in mylist:
    root = element ** (1 / 3)
    cube_roots.append(root)
    
    print(f"The cube root of {element} is {root:.4f}")

The cube root of 1 is 1.0000
The cube root of 2 is 1.2599
The cube root of 3 is 1.4422
The cube root of 4 is 1.5874
The cube root of 5 is 1.7100
The cube root of 6 is 1.8171
The cube root of 7 is 1.9129

But this is somewhat of a lot of effort just to create this new list. We could of course make use of numpy’s abilities to do this more simply with cube_roots = np.power(mylist, 1/3), but sometimes using an array is not possible, or we need to use a list. This is where the list comprehension allows us to compact this loop into one line:

mylist = [1, 2, 3, 4, 5, 6, 7]
cube_roots = [el ** (1 / 3) for el in mylist]
print(cube_roots)

[1.0, 1.2599210498948732, 1.4422495703074083, 1.5874010519681994, 1.7099759466766968, 1.8171205928321397, 1.912931182772389]

To unpack this, we can understand the basic syntax as

[*instructions* for *element* in *iterable*]

So in the example above, the instructions were el ** (1 / 3), the element was a variable we called el and the iterable was mylist. These instructions can typically be of any complexity as long as they don’t involve Python’s multiline syntax - you can’t insert a regular for loop, for example. The elements and iterables can be just as complicated as in regular loops, but the basic syntax will remain the same. Consider for example:

mylist = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
newlist = [el ** (1 / (idx + 1)) for (idx, el) in enumerate(mylist)]
print(newlist)

[1.0, 1.4142135623730951, 1.4422495703074083, 1.4142135623730951, 1.379729661461215, 1.3480061545972777, 1.3204692477561237, 1.2968395546510096, 1.2765180070092417, 1.2589254117941673]

Here we have created a list that contains the elements of mylist each raised to the power of \(1 / (`idx` + 1)\), where idx is the index location of the element in mylist. That is, this example made use of the enumerate function that we’ve previously seen in the section on loops, and we were then able to create a variable for both the element and its index within the list comprehension.

There are more modifications that can be made to list comprehensions, but this should serve as a sufficient introduction. For more information, I recommend this page for some other simple examples. Otherwise, feel free to continue on to the section on conditional expressions.