Python List Comprehension

Python list comprehension allows us to create and manipulate lists in a much easier way. Before we learn about list comprehension, let's see how we create lists normally.

# create a list of numbers
even_numbers = [1, 2, 3, 4]
print(even_numbers)

# Output: [1, 2, 3, 4]

Now, let's create this same list using list comprehension.

# create a list of numbers using list comprehension
numbers = [num for num in range(1, 5)]

print(numbers)

# Output: [1, 2, 3, 4]

From the above example, we can derive the syntax of list comprehension.

Syntax of List Comprehension

[expression for element in sequence]

List Comprehension to Create a List from Existing List

We can also use list comprehension to create a new list from the existing list. For example,

Suppose we have a list of numbers, and we are creating a new list by squaring each element of the list.

# a list of numbers
numbers = [1, 2, 3, 4, 5]
print('Old List:', numbers)

# create a new list using list comprehension
square_numbers = [num ** 2 for num in numbers]

print('New List:', square_numbers)

Output

Old List: [1, 2, 3, 4, 5]
New List: [1, 4, 9, 16, 25]

In the above example, we have used the list comprehension to create a new list where each element is the square value of the element of the existing list.

square_numbers = [num * num for num in numbers]

Here, each value of the numbers list is assigned to the num variable, and after performing num * num, the result is stored in the new list.

for Loop Vs. List Comprehension

In the previous example, we used list comprehension to create a new list by squaring each value of the existing list.

Now, let's perform the same operation using a for loop.

# a list of numbers
numbers = [1, 2, 3, 4, 5]
print('Old List:', numbers)

# a new empty list
square_numbers = []

for num in numbers:
    square_numbers.append(num**2)
    
print('New List:', square_numbers)

Output

Old List: [1, 2, 3, 4, 5]
New List: [1, 4, 9, 16, 25]

When we compare the above code with the one using list comprehension, it's clear that the use of list comprehension makes our code cleaner and easier to understand.

Function Call from List Comprehension

We can also call a function from within the list comprehension. For example,

# function to find the square of a number
def find_square(number):
    return number**2
    
# a list of numbers
numbers = [5, 7, 9, 11]

# use list comprehension to find the square
square_numbers = [find_square(num) for num in numbers]

print('Square Numbers:', square_numbers)

Output

Square Numbers: [25, 49, 81, 121]

In the above example, notice that we are calling the find_square(num) function from within the list comprehension.

square_numbers = [find_square(num) for num in numbers]

Here, for each value in the numbers list, the find_square() function is called, and a new list is created based on the returned value.

if…else with List Comprehension

We can also include a conditional statement like if else as an expression of the list comprehension. For example,

# create a list of even numbers
even_numbers = [num for num in range(1, 10) if num % 2 == 0 ]

print(even_numbers)

# Output: [2, 4, 6, 8]

In the above example, for every number in range 1 to 10 (excluding 10), the condition num % 2 == 0 is applied.

If the condition is True, then the value will be added to the new list. That's why we get a list of only even numbers.

We can also include multiple conditions with a list comprehension. For example,

# create a list of negative even numbers
even_numbers = [num for num in range(-10, 10) if num < 0 if num % 2 == 0 ]

print(even_numbers)

# Output: [-10, -8, -6, -4, -2]

Here, we have used two if conditions: if num < 0 and num % 2 == 0. In this case, only those elements that satisfy both conditions are added to the list.

That's why we get only negative even numbers.

Now let's see one more example of using the else statement as well.

# list of person ages
ages = [21, 16, 17, 23, 18]

# list to check if person can vote
can_vote = ['Yes' if age >= 18 else 'No' for age in ages]

print(can_vote)

# Output: ['Yes', 'No', 'No', 'Yes', 'Yes']

In this example, we are given a list of ages. From the list, we are finding the number of people who can vote.

Here, if the age is greater than or equal to 18, Yes is assigned to the new list, else, No is assigned.

Nested List Comprehension

We can also create a nested list using list comprehension. For example,

# create a nested list of multiplication table
multiplication = [[i * j for j in range(1, 6)] for i in range(2, 5)]
print(multiplication)

Output

[[2, 4, 6, 8, 10], [3, 6, 9, 12, 15], [4, 8, 12, 16, 20]]

Here, you can see we have successfully created a nested list. Let's see how this program works.

• i * j is the expression of for j in range(1, 6)
• [i * j for j in range(1, 6)] is the expression of the for i in range(2, 5)]
• for each element in range 2 to 5 (excluding 5), the inner list comprehension is executed
• and for each element in range 1 to 6 (excluding 6), i * j is executed.

List Comprehension and Lambda Function

In Python, we can also use lambda functions to perform similar operations like list comprehension. Let's see one example.

Here's the code to square all values of a list using list comprehension.

# a list of numbers
numbers = [5, 6, 7, 8, 9]
print('Old List:', numbers)

# create a new list using list comprehension
square_numbers = [num ** 2 for num in numbers]

print('New List:', square_numbers)

Output

Old List: [5, 6, 7, 8, 9]
New List: [25, 36, 49, 64, 81]

Now let's do the same thing using a lambda function.

# a list of numbers
numbers = [5, 6, 7, 8, 9]
print('Old List:', numbers)

# create a new list using a lambda function
square_numbers = list(map(lambda num : num**2 , numbers))
print('New List:', square_numbers)

Output

Old List: [5, 6, 7, 8, 9]
New List: [25, 36, 49, 64, 81]

Here, you can see we have successfully achieved the same result using a lambda function. To learn more, visit Python Lambda Function.

Benefits of List Comprehension

Here're some of the major advantages of using list comprehension in Python.

• List comprehension is closer to English, so it's easier to read and understand than the long Python code.
• As we have seen above, list comprehension makes working with python a lot easier.
• It is considered the most Pythonic way to write code as it is such a simple yet powerful tool.
• It is considered to be more efficient than the for loop while working with lists in Python.