How to use list comprehension in Python

Jan 05, 2024#python#lists

List comprehension is a concise way of creating lists in Python based on an existing iterable object, such as a list, a tuple, a string, or a range. It allows you to apply an expression or a condition to each element and produce a new list as a result.

new_list = [expression for element in iterable if condition]

"""
- new_list: name of the list you want to create.
- expression: what you want to do with each element, such as a mathematical operation, a function call, or a string manipulation.
- element: a variable that represents each item in the iterable object.
- iterable: the object that you want to iterate over, such as a list, a tuple, a string, or a range.
- condition: an optional filter that only includes the elements that satisfy a certain criterion, such as a comparison, a membership test, or a logical operator.
"""

List comprehension is generally faster than for loops, especially for large datasets. It takes less code to write and fits in a smaller space than a for loop. It can also be used to create other types of collections, such as sets and dictionaries.

  1. Creating new lists based on existing iterables
squares = [x**2 for x in range(10)]
# [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

words = [word for word in "This is a sentence with words.".split()]  
# ["This", "is", "a", "sentence", "with", "words"]
  1. Filtering elements based on conditions
numbers = [1, 2, 3, 4, 5, 6]
even_numbers = [num for num in numbers if num % 2 == 0]  
# [2, 4, 6]

strings = ["apple", "banana", "cherry", "date"]
long_strings = [str for str in strings if len(str) > 5]  
# ["banana", "cherry"]
  1. Applying transformations to elements
names = ["john", "jane", "doe"]
uppercase_names = [name.upper() for name in names]  
# ["JOHN", "JANE", "DOE"]

numbers = [1, 2, 3, 4]
tripled_numbers = [num * 3 for num in numbers]  
# [3, 6, 9, 12]
  1. Flatten a nested list or a matrix into a single list
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flat_list = [x for row in matrix for x in row]
# [1, 2, 3, 4, 5, 6, 7, 8, 9]
  1. Combining multiple operations in a single expression
words = ["apple", "banana", "cherry"]
word_lengths = [(word, len(word)) for word in words]  
# [("apple", 5), ("banana", 6), ("cherry", 6)]

numbers = [1, 4, 9, 16, 25]
even_square_roots = [num**0.5 for num in numbers if num % 2 == 0]  
# [2.0, 4.0]

Set and dictionary comprehension

Set and dictionary comprehension in Python are similar to list comprehension, but they use curly braces instead of square brackets. They allow you to create sets and dictionaries based on existing iterables, such as lists, tuples, strings, sets, dictionaries, etc.

Set comprehension creates a set, which is an unordered collection of unique elements. Dictionary comprehension creates a dictionary, which is an ordered collection of key-value pairs.

new_set = {expression for item in iterable if condition}
new_dict = {key-expression: value-expression for item in iterable if condition}

Here’s an example that creates a set of odd numbers from 0 to 10:

odd_numbers = {x for x in range(11) if x % 2 == 1}
# {1, 3, 5, 7, 9}

Here’s an example that creates a dictionary that maps each letter in a word to its index:

word = "hello"
letter_index = {letter: index for index, letter in enumerate(word)}
# {'h': 0, 'e': 1, 'l': 3, 'o': 4}

Alternatives to list comprehension

There are several alternatives to list comprehension in Python, depending on your use case and preference. Each of these alternatives has its own advantages and disadvantages.

For loops: This is the most basic and intuitive way to create lists by iterating over an existing iterable and appending elements to a new list.

squares = []
for x in range(10):
    squares.append(x**2)

print(squares)
# [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Map function: This is a functional programming approach that applies a given function to each element of an iterable and returns a map object. You can convert the map object to a list using the list function.

def square(x):
    return x**2

squares = list(map(square, range(10)))

print(squares)
# [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Generator expressions: This is a similar syntax to list comprehension, but instead of creating a list, it creates a generator object that can be iterated over lazily. This can save memory and improve performance for large datasets. You can use parentheses instead of brackets to create a generator expression.

squares = (x**2 for x in range(10))

for value in squares:
    print(value)

# 0
# 1
# 4
# 9
# 16
# 25
# 36
# 49
# 64
# 81