note.nkmk.me

List comprehensions in Python

Posted: 2020-10-21 / Tags: Python, List

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

This post 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
  • Dict comprehensions
  • Generator expression

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

See the following post for examples of using list comprehensions.

Sponsored Link

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]

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]

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 expression part of list comprehensions as follows.

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 operation 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 for statement, zip() that aggregates elements from multiple iterables and enumerate() that returns a value with an index are often used.

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]
Sponsored Link

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 as a grammar.

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

Dict comprehensions

Dictionaries (dict type objects) can also be created with dict comprehensions.

Use {} and specify the key and the value in 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}

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 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 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 is 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'>
Sponsored Link
Share

Related Categories

Related Posts