Understand List comprehension in Python

·

5 min read

list comprehension in Python-original article

List comprehension is a type of syntax for quickly and conveniently creating lists and other sequences from a collection of elements.

The syntax combines the features of loops, conditional execution and sequence building all into one concise syntax.

Its basic syntax is as shown below:

Syntax:
[<expression> for <item> in <collection>]

The single arbitrary expression gets evaluated for items in the collections and the resulting value is added to the list.

The following example uses list comprehension to create a list from a range.

Example
L = [i for i in range(10)]

print(L)
Output:

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Literally speaking there is no much operational difference between the above example and manually using a for loop to append elements to the list as in the following example.

Example
L = []

for i in range(10):
    L.append(i)

print(L)
Output:

[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

As a matter of fact, this is what actually goes on internally when using list comprehension. However, the list comprehension syntax offers a more concise and readable way of expressing the same logic.

We can use arbitrary single expression in the first part of the syntax.

Example
languages = ('python', 'java', 'php', 'ruby')

L = [i.upper() for i in languages]

print(L)
Output:

['PYTHON', 'JAVA', 'PHP', 'RUBY']

In the above example, we used the upper() method of string objects so that each element is first transformed to uppercase before being appended to the list.

As in traditional for loops, if the collection is made up of two-dimensional iterables of uniform known length, each inner element can be aliased with a different name.

Example
T = ((1, 2), (3, 4), (5, 6))

L = [a + b for a, b in T]

print(L)
Output:

[3, 7, 11]

Extended syntax

List comprehension syntax can be made to perform more complex by including an if clause. This way, the syntax gets extended as shown below.

Syntax:
[<expression> for <item> in <collection> if <condition>]

When the if clause is used, the expression gets executed only if the condition is met.

Example
L = [i for i in range(20) if i % 2 == 0]

print(L)
Output:

[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]

The above expression creates a list of numbers from 0 to 19 that are divisible by 2.

Example
T = ('JAVA', 'c++', 'php', 'PYTHON', 'RUBY')

L = [i for i in T if i.isupper()]

print(L)
Output:

['JAVA', 'PYTHON', 'RUBY']

In the above example we used the isupper() method of strings to specify the condition so that only string whose all letters are capital will be included in the resulting list.

Chaining multiple conditions

We can specify more than one condition, in which case they will be evaluated together.

Example
L = [x for x in range(20) if x%2 == 0 if x%3 == 0]

print(L)
Output:

[0, 6, 12, 18]

This above example results in a list with all numbers between 0 and 20 that are divisible by both 2 and 3.

Multiple for loops in a list comprehension

We can use more than one for loop in the list comprehension syntax, this is especially useful if we are dealing with a nested iterable. Consider the following example.

Example
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

#flatten the matrix
result = [item for row in matrix for item in row]

print(result)
Output:

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

In the above example in line result = [item for row in matrix for item in row] ,:

  • The first for loop iterates over the items in the matrix, it takes each row in the matrix and stores it in the variable 'row'.

  • The second for loop iterates through each element of the row and stores it in the variable 'item'

  • Each item is appended to the list during each iteration.

  • The net outcome is that the resulting list contains the individual elements of the inner lists.

Nested List comprehension syntax

The expression in a list comprehension can be any valid single Python expression. This means that we can also have a list comprehension inside another list comprehension. This type of list comprehension is known as a nested list comprehension.

A nested list comprehension can be used to conveniently iterate through an iterable of iterables (such as a list of lists ). This allows us to create a list containing all the elements of all the iterables. Consider the following example:

Example
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] 

#List comprehension to square every item in the matrix 
result = [[item ** 2 for item in row] for row in matrix]

print(result)
Output:

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

In the above example, the first for loop is going through each item in each row and squares them, while the second for loop goes through each row in the matrix.

The above list comprehension example is roughly equivalent to the example below which uses for loops:

Example
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] 

result = []

for row in matrix:
    temp_list = []
    for item in row:
       temp_list.append(item ** 2)

    result.append(temp_list)

print(result)
Output:

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