List comprehensions in Python

Modified: | Tags: Python, List

In Python, you can create a new list using list comprehensions. It's simpler than using the for loop.

See the following article for the basics of the for statement.

See the following article for examples of using list comprehensions.

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 assigned to 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 provided alongside an equivalent for statement. Squares the elements of range.

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 satisfy the condition are evaluated with expression. Elements that do not meet the condition are excluded from the resulting list.

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 previous example, elements that do not satisfy the 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 the value or expression when the condition is True, and Y is the value or expression when the condition is 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 have been added for readability, 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}
source: set.py

For more information about set, see the following article.

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.

{key: value for variable_name in iterable}

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}

For other ways to create dictionaries, see the following articles

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 generate elements one at a time during iteration, which can reduce memory usage.

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

When all elements are processed, list comprehensions are often faster than generator expressions.

However, in all() and any(), for example, a generator expression may be faster than list comprehensions because the result is determined when False or True is encountered.

There are no tuple comprehensions, but it is possible to create tuples by passing the generator expression to tuple().

t = tuple(i**2 for i in range(5))

print(t)
# (0, 1, 4, 9, 16)

print(type(t))
# <class 'tuple'>

Related Categories

Related Articles