Reprinted from: note.nkmk.me-List comprehensions in Python
In Python, you can create a new list using list comprehensions. It's simpler than using the for loop.
This article describes the following contents.
- Basics of list comprehensions
- List comprehensions with
if - List comprehensions with conditional expressions (like
if else) - List comprehensions with
zip(),enumerate() - Nested list comprehensions
- Set comprehensions
- Dictionary comprehensions
- Generator expression
See the following article for examples of using list comprehensions.
- Extract, replace, convert elements of a list in Python
- Extract and replace elements that meet the conditions of a list of strings in Python
Basics of list comprehensions
List comprehensions are written as follows:
[expression for variable_name in iterable]
Each element of iterable, such as a list or a tuple, is taken out as variable_name and evaluated with expression. A new list is created with the result evaluated by expression as an element.
An example of list comprehensions is shown with an equivalent for statement.
squares = [i**2 for i in range(5)]
print(squares)
# [0, 1, 4, 9, 16]
squares = []
for i in range(5):
squares.append(i**2)
print(squares)
# [0, 1, 4, 9, 16]
Although map() can do the same thing, list comprehension is preferred for simplicity and clarity.
List comprehensions with if
In list comprehensions, if can be used as follows:
[expression for variable_name in iterable if condition]
Only elements of iterable that are True for condition are evaluated with expression. Elements of iterable that are False for condition are not included in the list of results.
odds = [i for i in range(10) if i % 2 == 1]
print(odds)
# [1, 3, 5, 7, 9]
odds = []
for i in range(10):
if i % 2 == 1:
odds.append(i)
print(odds)
# [1, 3, 5, 7, 9]
Although filter() can do the same thing, list comprehension is preferred for simplicity and clarity.
List comprehensions with conditional expressions (like if else)
In the above example, elements that do not meet condition are excluded from the new list.
If you want to apply another operation to elements that do not meet condition like if else, use conditional expressions.
In Python, conditional expressions can be written as follows:
X if condition else Y
X is value or expression for True, and Y is value or expression for False.
You can use it for the expression part of list comprehensions:
odd_even = ['odd' if i % 2 == 1 else 'even' for i in range(10)]
print(odd_even)
# ['even', 'odd', 'even', 'odd', 'even', 'odd', 'even', 'odd', 'even', 'odd']
odd_even = []
for i in range(10):
if i % 2 == 1:
odd_even.append('odd')
else:
odd_even.append('even')
print(odd_even)
# ['even', 'odd', 'even', 'odd', 'even', 'odd', 'even', 'odd', 'even', 'odd']
It is also possible to apply different operations to the original element depending on the condition as follows:
odd10 = [i * 10 if i % 2 == 1 else i for i in range(10)]
print(odd10)
# [0, 10, 2, 30, 4, 50, 6, 70, 8, 90]
List comprehensions with zip(), enumerate()
In the for statement, zip() and enumerate() are often used. zip() aggregates elements from multiple iterables and enumerate() returns a value and its index.
Of course, you can also use zip() and enumerate() in list comprehensions.
List comprehensions with zip():
l_str1 = ['a', 'b', 'c']
l_str2 = ['x', 'y', 'z']
l_zip = [(s1, s2) for s1, s2 in zip(l_str1, l_str2)]
print(l_zip)
# [('a', 'x'), ('b', 'y'), ('c', 'z')]
l_zip = []
for s1, s2 in zip(l_str1, l_str2):
l_zip.append((s1, s2))
print(l_zip)
# [('a', 'x'), ('b', 'y'), ('c', 'z')]
List comprehensions with enumerate():
l_enu = [(i, s) for i, s in enumerate(l_str1)]
print(l_enu)
# [(0, 'a'), (1, 'b'), (2, 'c')]
l_enu = []
for i, s in enumerate(l_str1):
l_enu.append((i, s))
print(l_enu)
# [(0, 'a'), (1, 'b'), (2, 'c')]
You can also use if.
l_zip_if = [(s1, s2) for s1, s2 in zip(l_str1, l_str2) if s1 != 'b']
print(l_zip_if)
# [('a', 'x'), ('c', 'z')]
It is also possible to calculate a new value using each element.
l_int1 = [1, 2, 3]
l_int2 = [10, 20, 30]
l_sub = [i2 - i1 for i1, i2 in zip(l_int1, l_int2)]
print(l_sub)
[9, 18, 27]
Nested list comprehensions
List comprehensions can be nested just as for loops are nested.
[expression for variable_name1 in iterable1
for variable_name2 in iterable2
for variable_name3 in iterable3 ... ]
Line breaks and indentation are added for convenience, but they are not required.
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flat = [x for row in matrix for x in row]
print(flat)
# [1, 2, 3, 4, 5, 6, 7, 8, 9]
flat = []
for row in matrix:
for x in row:
flat.append(x)
print(flat)
# [1, 2, 3, 4, 5, 6, 7, 8, 9]
Another example:
cells = [(row, col) for row in range(3) for col in range(2)]
print(cells)
# [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)]
You can also use if.
cells = [(row, col) for row in range(3)
for col in range(2) if col == row]
print(cells)
# [(0, 0), (1, 1)]
cells = [(row, col) for row in range(3) if row % 2 == 0
for col in range(2) if col % 2 == 0]
print(cells)
# [(0, 0), (2, 0)]
Set comprehensions
If you change the square brackets [] in the list comprehension to the curly braces {}, a set (set type object) is created.
{expression for variable_name in iterable}
s = {i**2 for i in range(5)}
print(s)
# {0, 1, 4, 9, 16}
Dictionary comprehensions
Dictionaries (dict type objects) can also be created with dictionary comprehensions.
Use {} and specify the key and the value in the expression part as key: value.
Any value or expression can be specified for key and value.
l = ['Alice', 'Bob', 'Charlie']
d = {s: len(s) for s in l}
print(d)
# {'Alice': 5, 'Bob': 3, 'Charlie': 7}
Use the zip() function to create a new dictionary from each list of keys and values.
keys = ['k1', 'k2', 'k3']
values = [1, 2, 3]
d = {k: v for k, v in zip(keys, values)}
print(d)
# {'k1': 1, 'k2': 2, 'k3': 3}
Generator expressions
If the square bracket [] in the list comprehension is changed to the parenthesis (), the generator is returned instead of the tuple. This is called a generator expression.
List comprehension:
l = [i**2 for i in range(5)]
print(l)
# [0, 1, 4, 9, 16]
print(type(l))
# <class 'list'>
Generator expression:
print() does not output the generator elements. You can get elements by using the for statement.
g = (i**2 for i in range(5))
print(g)
# <generator object <genexpr> at 0x10af944f8>
print(type(g))
# <class 'generator'>
for i in g:
print(i)
# 0
# 1
# 4
# 9
# 16
It is possible to use if and nest with generator expressions as well as list comprehensions.
g_cells = ((row, col) for row in range(0, 3)
for col in range(0, 2) if col == row)
print(type(g_cells))
# <class 'generator'>
for i in g_cells:
print(i)
# (0, 0)
# (1, 1)
For example, if a list with many elements is created with list comprehensions and used in the for loop, list comprehensions create a list containing all elements at first.
On the other hand, generator expressions create elements one by one after each loop, so the memory usage can be reduced.
You can omit the parentheses () if the generator expression is the only argument passed to the function.
print(sum([i**2 for i in range(5)]))
# 30
print(sum((i**2 for i in range(5))))
# 30
print(sum(i**2 for i in range(5)))
# 30
```
There are no tuple comprehensions, but it is possible to create tuples by passing the generator expression to `tuple()`.
```python
t = tuple(i**2 for i in range(5))
print(t)
# (0, 1, 4, 9, 16)
print(type(t))
# <class 'tuple'>
```