Replace Loops, map() and filter() With List Comprehensions in Python

I’m a fan of The Zen of Python, which says

There should be one — and preferably only one — obvious way to do it.

But in Python, there are in fact many ways to achieve the same objective. Of course, some ways are more elegant than others and in most cases, it should be obvious which way is better.

We are going to look at list comprehensions, and how they can replace for loops, map() and filter() to create powerful functionality within a single line of Python code.

Basic List Comprehension

Say I want to create a list of numbers from 1 to 10. I could do

numbers = []
for i in range(1, 11):
    numbers.append(i)

and I would get

>>> numbers
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

But using list comprehension, this could be accomplished within a single line

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

This is the basic syntax of list comprehension: [expression for element in iterable]. Here, the iterable is range(1, 11), the element is i and the expression is i. This is equivalent to the for loop we used earlier: we add i to the list where i is a number from 1 to 11.

map()

The map() function is often used to apply a function on each element in an iterable. Pass in a function and an iterable, and map() will create an object containing the results of passing each element into the function.

For example, say I want to create a list of squares from the numbers list we created earlier. We could do

squares = []
for num in numbers:
    squares.append(num ** 2)

and we will get

>>> squares
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

Or we could use map() like so

>>> squares = list(map(lambda x: x ** 2, numbers))
>>> squares
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

Here, we pass each of the elements in numbers into the lambda function (which is just an easy way to create functions if you are lazy). The output of putting each number x into the function lambda x: x ** 2will be the square of the number. Using list() we turn the map object into a list.

Replacing map() With List Comprehension

Using list comprehension, we could simply do

>>> squares = [num ** 2 for num in numbers]
>>> squares
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

This will pass each number num into the expression num ** 2 and create a new list where the elements are simply the squares of each number in numbers.

filter()

The filter() function is used to create a subset of an existing list, based on a condition. Pass in a function and an iterable, and filter() will create an object containing all elements where the function evaluates to True.

For example, say I want to get a list of all odd numbers and a list of all even numbers from the numbers list we created earlier. We could do

odd_numbers = []
even_numbers = []

for num in numbers:

    if num % 2 == 1:
        odd_numbers.append(num)

    elif num % 2 == 0:
        even_numbers.append(num)

and we would get

>>> odd_numbers
[1, 3, 5, 7, 9]
>>> even_numbers
[2, 4, 6, 8, 10]

Or we could use filter() like so

odd_numbers = list(filter(lambda x: x % 2 == 1, numbers))
even_numbers = list(filter(lambda x: x % 2 == 0, numbers))

Here, all the numbers will be passed into the lambda function and if x % 2 == 1 is True, the number will be included in odd_numbers. Likewise, if x % 2 == 0 is True, the number will be included in even_numbers.

Replacing filter() With List Comprehension

Using list comprehension, we could simply do

odd_numbers = [num for num in numbers if num % 2 == 1]
even_numbers = [num for num in numbers if num % 2 == 0]

Here, we are using a conditional. The syntax for this is [expression for element in iterable (if conditional)]. This is equivalent to the for loop we used earlier — if the condition num % 2 == 1 is met, num will be added to odd_numbers, and num is an element in the numbers iterable.

Nested Comprehensions

Say we want to create a matrix. This would involve creating nested lists. Using a for loop, we can do the following

matrix = []
for i in range(5):
    row = []
    for j in range(5):
        row.append(i * 5 + j)
    matrix.append(row)

and we would get

>>> matrix
[
    [0, 1, 2, 3, 4],
    [5, 6, 7, 8, 9],
    [10, 11, 12, 13, 14],
    [15, 16, 17, 18, 19],
    [20, 21, 22, 23, 24]
]

Using a nested comprehension, we could simply do

matrix = [[i * 5 + j for j in range(5)] for i in range(5)]

The outer list comprehension [... for i in range(5)] creates 5 rows, while the inner list comprehension [... for j in range(5)] creates 5 columns.

Dictionary Comprehensions

You could use dictionary comprehensions too. For example, if I want to create a dictionary that maps each number in numbers to their corresponding square, we could do

num_to_square = {}
for num in numbers:
    num_to_square[num] = num ** 2

and we would get

>>> num_to_square
{
    1: 1,
    2: 4,
    3: 9,
    4: 16,
    5: 25,
    6: 36,
    7: 49,
    8: 64,
    9: 81,
    10: 100
}

Or we could use a dictionary comprehension

>>> num_to_square = {num: num ** 2 for num in numbers}
>>> num_to_square
{
    1: 1,
    2: 4,
    3: 9,
    4: 16,
    5: 25,
    6: 36,
    7: 49,
    8: 64, 
    9: 81,
    10: 100
}

You could even use a dictionary comprehension within a list comprehension or vice versa.

>>> num_to_square = [{num: num ** 2 for num in numbers} for numbers in
[odd_numbers, even_numbers]]
>>> num_to_square
[
    {1: 1, 3: 9, 5: 25, 7: 49, 9: 81},
    {2: 4, 4: 16, 6: 36, 8: 64, 10: 100}
]

Why Use List Comprehension?

• Cleaner, clearer code
• Slightly faster than map() and filter()
• Generally considered more ‘pythonic’

But hey, at the end of the day, the choice is yours. List comprehension is just another way to do the same things, and whether something is ‘cleaner’ and ‘clearer’ is largely subjective. However, most people would agree that list comprehension is a more standard way to do things.

Conclusion

That’s it! Now you know how to use list comprehensions in Python.

If you have any questions, feel free to comment down below.

Thanks for reading, and follow me for more stories like this in the future.