Understanding List Comprehension in Python

Proposed to allow the concise construction of lists, list comprehension offers Python developers an elegant alternate solution to nesting multiple FOR loops and IF clauses. The concise nature of list comprehension can greatly reduce the space used by the code. Therefore, the visual noise is reduced, and the overall readability of the code is improved when list comprehension is used appropriately.

In Python, list comprehension takes place within the list syntax and is generally constructed in the following form:

[output iterable condition]

The output is listed first, next an iterable is identified (such as a list), and last any optional conditions can be listed.

Let’s take a look at list comprehension in a basic form. Below is a Python code snippet that creates a list using numbers 0 to 9.

number_list = []
for i in range(10):
>>> [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

With List Comprehension, this can be shortened to the following:

number_list = [i for i in range(10)]

At first, this may seem confusing; however, deconstructing it into smaller parts should clarify the code. The left-most variable is the output; it’s equivalent to the append(i) in the previous code snippet. The for i section defines a variable i that iterates over an iterable range(10).

Although this is rather basic, it demonstrates how concise list comprehension can make your code thus reducing overall visual noise and improving the general readability. Even though list comprehension can seem complex at first glance, it generally gives the visual impression of being an isolated section of code that can be analysed on its own without worrying about the surrounding sections of code.

Output Modification

What if we want a different output from the iterable? For example, we’re trying to generate a list of square values from the code. A standard way of coding this would be to write:

squares = []
for x in range(10):
>>> [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

This too can be simplified using list comprehension with a small modification to the above code:

squares = [x**2 for x in range(10)]

Notice that the output variable (the first statement) is modified to obtain the squared result. Again this is equivalent to the append(x**2) of the previous code.

Changing the Input

Any iterable can be used with list comprehension (list, str, tuple); even non-sequence types that expose the iter() function can be used (dict, object). Here are some example inputs:

# Convert the string 'split' into a list where each element is a character
[x for x in 'split']
>>> ['s', 'p', 'l', 'i', 't']

# Convert a tuple into a list
[x for x in (10, 20, 30)]
>>> [10, 20, 30]

# Extract the first element of each tuple and multiply it by 2
[x[0]*2 for x in [(1, 2), (3, 4), (5, 6)]]
>>> [2, 6, 10]

Conditions can be applied as well. The following example demonstrates a dictionary output of squares only if the input value is greater than 5.

[{'x': x, 'x**2': x**2} for x in range(5) if x > 2]
>>> [{'x': 3, 'x**2': 9}, {'x': 4, 'x**2': 16}]

What if we want to modify the output based on its value? We can do this too using an IF-ELSE clause on the left side. Here if x%2 evaluates to one (TRUE), then it’s odd, else its even.

[str(x) + ': Odd' if x % 2 else str(x) + ': Even' for x in range(5)]
>>> ['0: Even', '1: Odd', '2: Even', '3: Odd', '4: Even']
Nested Loops

Here is where list comprehension either becomes the saviour of readable and concise code or the bane of your project. If list comprehensions are used in overly complex situations, they can create an unintelligible tangle of code that will be difficult to debug or decipher by other developers.

To demonstrate a situation for nested loops, let’s consider a situation where we’re trying to find all the cartesian integer coordinates within a circle of radius 2 that is centred around the (0, 0) point. For this, we’ll loop over the x and y coordinates and check if they’re within a radius of 2 using Pythagoras’ Theorem. Here we use range(-2, 3) which creates a list of [-2, -1, 0, 1, 2],

import math
radius = 2

points = []
for x in range(-2, 3):
    for y in range(-2, 3):
        r = math.sqrt(x**2 + y**2)
        if r < radius:
            points.append((x, y))
>>> [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 0), (0, 1), (1, -1), (1, 0), (1, 1)]

In list comprehension, this can be reduced using multiple FOR statements.

radius = 2
[(x, y) for x in range(-2, 3) for y in range(-2, 3) if math.sqrt(x**2 + y**2) < radius]

This can be extended to as many FOR loops as required.

Note that when nesting is too deep, trying to use list comprehension limits readability and generally has the opposite effect to what you want to achieve.


Related Posts