List Comprehension
List comprehensions are a highly interesting and useful Python feature that allow us to create and manipulate lists using a much shorter syntax.
Syntax
The basic syntax of a list comprehension consists of square brackets [] enclosing an expression followed by a for
clause. We can also include if
clauses for filtering.
list_comp = [expression for item in iterable if condition]
Let’s get into the details of these parts below.
Expression
The expression defines what to include in the new list for each item in the iterable. It can be a simple transformation or even a computation.
For instance, suppose we have a list of five numbers stored in numbers
. We want to create another list from it and store the squares of these numbers in it. We can do this by declaring a list comprehension called squares
that evaluates the square of each number through the expression.
numbers = [1, 2, 3, 4, 5]squares = [num ** 2 for num in numbers]print(squares)
Condition
We can use conditions to filter elements from the original iterable based on specified criteria.
Let’s filter a list containing numbers so that it only stores odd numbers. We can do this by adding an ‘if’ statement at the end of the list comprehension that only saves a number from the list if it is odd (i.e., num % 2 != 0).
numbers = [1, 2, 3, 4, 5]odd_numbers = [num for num in numbers if num % 2 != 0]print(odd_numbers)
Let’s take another example that only stores words if they are short (i.e., have a length of less than or equal to five).
words = ["hoodie", "rivers", "cat", "monitor", "bag", "cup"]short_words = [word for word in words if len(word) <= 5]print(short_words)
Assignment to multiple variables
List comprehensions can be used to assign values to multiple variables simultaneously. In the example below, we save both the number and its square as a pair in the pairs variables.
numbers = [1, 2, 3, 4, 5]pairs = [(num, num * 2) for num in numbers]print(pairs)
Transforming elements
We can also call functions to transform elements during the creation of the new list.
songs = ['Neon Lights', 'Pieces', 'Everything']uppercase_songs = [song.upper() for song in songs]print(uppercase_songs)
Nested list comprehension
Nested comprehensions allow us to process elements from nested iterables (i.e., a 2D array).
Let’s take a look at a more advanced example that takes the transpose of a matrix through two ‘for’ loops in a nested list comprehension.
my_matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]transpose = [[row[i] for row in my_matrix] for i in range(len(my_matrix[0]))]print(transpose)
Outer list comprehension (for i in range(len(my_matrix[0]))):
It iterates over the range of column indices in the original matrix.
For each column index i, it executes the inner list comprehension.
Inner list comprehension ([row[i] for row in my_matrix]):
For a specific column index i, it iterates over each row in the original matrix.
For each row, it extracts the element at the current column index i.
The result is a list containing all elements from the same column for the given index i.
The outer list comprehension as a whole creates a list of lists, where each inner list is the result of the inner comprehension for a specific column index. These inner lists collectively form the rows of the transposed matrix.