List Comprehension in Python

List comprehension provides a concise way to create a new list in a single line by applying an expression to each item of an existing iterable (such as a list, tuple, or range), with an optional condition for filtering items.

Basic Syntax

[expression for item in iterable]

Where,

expression: The operation to perform on each item.

item: Variable representing each item.

iterable: The sequence (list, tuple, or range) to iterate over.

Example: Squaring Numbers

Let’s say you have a list of numbers and you want to create a new list with the squares of each number. Here is how you can do it using both a traditional for loop and list comprehension.

Using a for loop

numbers = [1, 2, 3, 4]
squares = []

for num in numbers:
    squares.append(num ** 2)
    
print(squares) # Output: [1, 4, 9, 16]

Using List Comprehension

numbers = [1, 2, 3, 4]
squares = [num ** 2 for num in numbers]
print(squares) # Output: [1, 4, 9, 16]

As you can see, list comprehensions achieve the same result as for loop, but with shorter and cleaner code.

Example: Create a List of Squares

squares = [num ** 2 for num in range(10)]
print(squares) # Output: [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

The range() generates a sequence of numbers. The expression range(10) generates numbers from 0 up to (but not including) 10.

Example: Convert List of Strings to Uppercase

fruits = ["apple", "banana", "mango", "orange"]

fruits_upper = [fruit.upper() for fruit in fruits]
print(fruits_upper) # Output: ['APPLE', 'BANANA', 'MANGO', 'ORANGE']

Conditional Logic With List Comprehension

You can add an if statement at the end of a list comprehension to include only the items that meet a certain condition from the original iterable. Here is a syntax of list comprehension with conditional logic:

[expression for item in iterable if condition]

Example: Getting Only Even Numbers

Let’s filter a list to get only the even numbers.

numbers = [1, 2, 3, 4, 5, 6, 7]
even_numbers = [num for num in numbers if num % 2 == 0]
print(even_numbers)

Numbers that are completely divisible by 2 , meaning num % 2 == 0 are even numbers.

if-else in List Comprehension

You can also use if-else within the expression part of the list comprehension. The syntax for list comprehension with if-else is below:

[expression_if_true if condition else expression_if_false for item in iterable]

Example: Odd or Even

Let’s use a list comprehension with an if-else condition to classify each number in a list as odd or even.

numbers = [1, 2, 3, 4, 5]
labels = ["Even" if num % 2 == 0 else "Odd" for num in numbers]
print(labels) # Output: ['Odd', 'Even', 'Odd', 'Even', 'Odd']

Nested List Comprehension

A nested list comprehension is a list comprehension inside another one. It helps simplify code when dealing with nested data, such as 2D lists, matrices, grids, or JSON data.

Example: Flattening a 2D List

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

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

The outer loop (for row in matrix) iterates through each row, and the inner loop (for num in row) extracts each number.

In a nested list comprehension, the outer loop starts first and picks one item at a time, and for each item from the outer loop, the inner loop runs completely before the outer loop moves to the next item.

Example: Create a Multiplication Table

table = [[i * j for j in range(1, 11)] for i in range(2, 5)]
print(table)

Output:

[
    [2, 4, 6, 8, 10, 12, 14, 16, 18, 20],  # 2 multiplication table
    [3, 6, 9, 12, 15, 18, 21, 24, 27, 30], # 3 multiplication table
    [4, 8, 12, 16, 20, 24, 28, 32, 36, 40] # 4 multiplication table
]

In this example, the outer loop (for i in range(2, 5)) goes through numbers 2, 3, and 4. These represent the multiplication table you want to create.

The inner loop (for j in range(1, 11)) goes through numbers 1 to 10 and multiplies each by i.

The result is a nested list where each inner list contains the multiplication table of 2, 3, and 4.