Python - List Comprehension

List comprehension in Python is an easy and compact syntax for creating a list from a string or another list. It is a very concise way to create a new list by performing an operation on each item in the existing list. List comprehension is considerably faster than processing a list using the for loop.

[expression for element in iterable if condition]

As per the above syntax, the list comprehension syntax contains three parts: an expression, one or more for loop, and optionally, one or more if conditions. The list comprehension must be in the square brackets []. The result of the first expression will be stored in the new list. The for loop is used to iterate over the iterable object that optionally includes the if condition.

Suppose we want to find even numbers from 0 to 20 then we can do it using a for loop, as shown below:

even_nums = []
for x in range(21):
    if x%2 == 0:
        even_nums.append(x)
print(even_nums)

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

The same result can be easily achieved using a list comprehension technique shown below.

even_nums = [x for x in range(21) if x%2 == 0]
print(even_nums)

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

In the above example, [x for x in range(21) if x%2 == 0] returns a new list using the list comprehension. First, it executes the for loop for x in range(21) if x%2 == 0. The element x would be returned if the specified condition if x%2 == 0 evaluates to True. If the condition evaluates to True, then the expression before for loop would be executed and stored in the new list. Here, expression x simply stores the value of x into a new list.

List comprehension works with string lists also. The following creates a new list of strings that contains 'a'.

names = ['Steve', 'Bill', 'Ram', 'Mohan', 'Abdul']
names2 = [s for s in names if 'a' in s]

print(names2)

['Ram', 'Mohan']

Above, the expression if 'a' in s returns True if an element contains a character 'a'. So, the new list will include names that contain 'a'.

The following example uses a list comprehension to build a list of squares of the numbers between 1 and 10.

squares = [x*x for x in range(11)] 
print(squares) 

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

Above, a for loop for x in range(11) is executed without any if condition. The expression before for loop x*x stores the square of the element in the new list.

List Comprehension using Nested Loops

It is possible to use nested loops in a list comprehension expression. In the following example, all combinations of items from two lists in the form of a tuple are added in a third list object.

nums1 = [1, 2, 3]
nums2 = [4, 5, 6]
nums=[(x,y) for x in nums1 for y in nums2]
print(nums)

[(1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6)]

List Comprehension with Multiple if Conditions

We can use nested if conditions with a list comprehension.

nums = [x for x in range(21) if x%2==0 if x%5==0] 
print(nums)

[0, 10, 20]

List Comprehension with if-else Condition

The following example demonstrates the if..else loop with a list comprehension.

odd_even_list = ["Even" if i%2==0 else "Odd" for i in range(5)]
print(odd_even_list)

odd_even_list = [str(i) + '=Even' if i%2==0 else str(i) + "=Odd" for i in range(5)]
print(odd_even_list)

['Even', 'Odd', 'Even', 'Odd', 'Even']
['0=Even', '1=Odd', '2=Even', '3=Odd', '4=Even']

Flatten List using List Comprehension

One of the applications of list comprehension is to flatten a list comprising of multiple lists into a single list.

matrix=[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flatList=[num for row in matrix for num in row]
print(flatList)

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

Learn more about how to flatten list in Python.