comprehensions.md

List Comprehensions

Source: Python List Comprehensions: Explained Visually by Trey Hunner

Conditional loop list comprehension

For loop:

numbers = [1, 2, 3, 4, 5]

doubled_odds = []
for n in numbers:
    if n % 2 == 1:
        doubled_odds.append(n * 2)

List comprehension:

numbers = [1, 2, 3, 4, 5]

doubled_odds = [n * 2 for n in numbers if n % 2 == 1]

Result

double_odds = [2, 6, 10]

Nested loop list comprehension

For loop:

matrix = [[1, 2], [3, 4], [5, 6]]

flattened = []
for row in matrix:
    for n in row:
        flattened.append(n)

List comprehension:

matrix = [[1, 2], [3, 4], [5, 6]]

flattened = [n for row in matrix for n in row]

Result

flattened = [1, 2, 3, 4, 5, 6]

Set comprehension

For loop:

words = ["Hi", "this", "is", "an", "example"]

first_letters = set()
for w in words:
    first_letters.add(w[0])

List comprehension:

words = ["Hi", "this", "is", "an", "example"]

first_letters = {w[0] for w in words}

Result

first_letters = {'H', 'i', 'e', 'a', 't'}

Dictionary comprehension

For loop:

original = {'k1':'v1', 'k2':'v2'}

flipped = {}
for key, value in original.items():
    flipped[value] = key

List comprehension:

original = {'k1':'v1', 'k2':'v2'}

flipped = {value: key for key, value in original.items()}

Result

flipped = {'v1':'k1', 'v2':'k2'}