List Comprehensions
Overview
List comprehensions let you create lists in a single line instead of using a for loop. It works with any iterable (lists, ranges, strings, etc) and shortens for loops.
Basic syntax:
[output_expression for item in iterable]
Comparison
List comprehensions work just like for loops but in a more compact way.
For example, consider the for
loop below:
nums = [1, 2, 3, 4]
new_nums = []
for num in nums:
new_nums.append(num + 1)
print(new_nums) ## Output: [2, 3, 4, 5]
This can rewritten as:
nums = [1, 2, 3, 4]
new_nums = [num + 1 for num in nums]
print(new_nums) ## Output: [2, 3, 4, 5]
Using range()
List comprehensions can work with range()
too.
squares = [x**2 for x in range(5)]
print(squares)
Output:
[0, 1, 4, 9, 16]
Nested Loops
You can use nested for loops inside a list comprehension.
Example: Creating Pairs
pairs = [(x, y) for x in range(2) for y in range(6, 8)]
print(pairs)
Output:
[(0, 6), (0, 7), (1, 6), (1, 7)]
Nested comprehensions can save space but may reduce readability. Use them when they make sense, but keep your code easy to understand.
Example: Matrices
Lists can store multi-dimensional data, like matrices. In Python, a matrix is just a list of lists. For example a 5 x 5 matrix with values 0 to 4 in each row can be written as:
matrix = [[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4],
[0, 1, 2, 3, 4]]
You can use a nested list comprehension to generate this matrix dynamically.
# 5 x 5 matrix using a list of lists
matrix = [[col for col in range(0, 5)] for row in range(0, 5)]
for row in matrix:
print(row)
Note that the output itself a list comprehension.
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]
[0, 1, 2, 3, 4]