Menu Close
  • Home
  • Python Comprehension Examples
    •

Python Comprehension Examples

Timon 0 Comments
comprehension

Introduction

Instead of writing out loops over multiple lines, the Python programming language has the possibility to use comprehensions, like the list comprehension. These are shortened ways of working with collections. The C# equivalent to this would be using LINQ, while for Java it would be using Streams. I did not research whether this goes against any design principle, I just show some possibilities. I used Python 3.13 in the examples.

List Comprehension

The syntax is the following:

newList = [ expression for i in iterable ]

For instance, to return a new list where the input is doubled:

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> List[int]:
        return [n*2 for n in numbers]


numbers = [1, 3, 5, 9, 0]
print(Example.comp(numbers)) # prints: [2, 6, 10, 18, 0]
Python

Without list comprehension, we could write this as:

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> List[int]:
        result = []
        for n in numbers:
            result.append(n * 2)
        return result


numbers = [1, 3, 5, 9, 0]
print(Example.comp(numbers))  # prints: [2, 6, 10, 18, 0]
Python

Two Dimensional Array

We could also iterate over multi dimensional arrays. This is called a nested list comprehension.

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[List[int]]) -> List[List[int]]:
        sums = [[n*2 for n in numList] for numList in numbers]
        return sums


numbers = [[1, 3, 5, 9, 0], [2, 2, 3], [7, 8, 2, 4]]
print(Example.comp(numbers)) # prints: [[2, 6, 10, 18, 0], [4, 4, 6], [14, 16, 4, 8]]
Python

Because our expression n*2 is inside two sets of [] it will also return a two dimensional array.

Let’s say we want to return the lowest sum of integers from multiple lists:

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[List[int]]) -> int:
        return min([sum(row) for row in numbers])

numbers = [[1, 3, 5, 9, 0], [2, 2, 3], [7, 8, 2, 4]]

print(Example.comp(numbers)) # prints: 7
Python

Iterate over all and print out all numbers in a one dimension array:

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[List[int]]) -> List[int]:
        return [n for row in numbers for n in row]


numbers = [[1, 3, 0], [2, 3], [7, 2, 4]]
print(Example.comp(numbers)) # prints: [1, 3, 0, 2, 3, 7, 2, 4]

# first row in numbers and first nested loop:
# for row in numbers will give [1, 3, 0] in the first iteration
# for n in row will give 1, 3, 0 sequentially in the nested loop
# it will append n to a list (behind the scenes)
Python

Transformed Comprehension

You can use if and else to apply a condition to each item. The syntax inside a list comprehension is as follows:

[ x if item condition else y for item in iterable ]

To apply a condition to each number:

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> List[str]:
        result = ["even" if n % 2 == 0 else "odd" for n in numbers]
        return result


numbers = [1, 3, 6, 9, 0]
print(Example.comp(numbers)) # prints: ['odd', 'odd', 'even', 'odd', 'even']
Python

We can also broaden the if-statement with an and:

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> List[int]:
        result = [n if n % 2 == 0 or n % 3 == 0 else -1 for n in numbers]
        return result


numbers = [1, 3, 6, 9, 11, 14]
print(Example.comp(numbers)) # prints: [-1, 3, 6, 9, -1, 14]
Python

Filtered Comprehension

Slighly different is the filter, which does not need an else. It just filters the loop from results which otherwise would go to the expression.

[ x for x in iterable if condition ]
from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> List[int]:
        result = [n for n in numbers if n % 2 == 0 or n % 3 == 0]
        return result
    
    
numbers = [1, 3, 6, 9, 11, 14]
print(Example.comp(numbers)) # prints: [3, 6, 9, 14]
Python

Because we did not have to write an else, the numbers that do not get through the filter are skipped.

Transformation and Filter Combined

“If” statements can be put on multiple positions and the position in the comprehension decides the role. An “if” before the for loop will transform the expression, behind it will filter it. In the next example, we can best start reading from the last loop, for n in numbers… then start reading from the beginning. So, if n is an even number, multiply by 2 if less than 5, divide by 2 if more than 5. Odd numbers get skipped.

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> List[int]:
        result = [int(n * 2) if n < 5 else int(n / 2) for n in numbers if n % 2 == 0]
        return result


numbers = [10, 3, 0, 2, 6, 5]
print(Example.comp(numbers)) # prints: [5, 0, 4, 3]
Python

Combine with Functions

from typing import List

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> List[str]:
        result = [Example.label(n) for n in numbers]
        return result

    def label(n):
        if n < 5:
            return "Low"
        elif n < 10:
            return "Medium"
        else:
            return "High"


numbers = [1, 3, 6, 9, 11, 14]
print(Example.comp(numbers)) 
# prints: ['Low', 'Low', 'Medium', 'Medium', 'High', 'High']
Python

Dictionary Comprehension

Dictionary comprehensions are a shortened way to work with dictionaries instead of lists. Basically, the same rules apply as for list comprehensions, with some small differences. A Dictionary has key/value pairs and we instantiate it with {} instead of [].

{ key : val*2 for val in iterable }
from typing import List, Dict

class Example:

    @staticmethod
    def comp(numbers: List[int]) -> Dict[int,int]:
        result = {index : val*2 for index, val in enumerate(numbers)}
        return result


numbers = [1, 2, 3]
print(Example.comp(numbers)) # prints: {0: 2, 1: 4, 2: 6}
Python

This also works with filters, transformations and dictionaries in dictionaries.

The next example will insert a list with users where each user has a Dictionary data type. The comp function will return a Dictionary with key : Dictionary for each active user, based on a filter, transformation and the use of a function.

from typing import List, Dict

class Example:

    def comp(self, users: List[Dict[str, any]]) -> Dict[int,Dict[str, str]]:
        user_dict = {
            user["id"]: {
                "name": user["name"].upper(),
                "status": "active" if user["active"] else "inactive",
                "grade": self.grade(user["score"])
            }
            for user in users
            if user["active"]
        }
        return user_dict

    def grade(self, score) -> str:
        if score < 50:
            return "F"
        elif score < 60:
            return "D"
        elif score < 70:
            return "C"
        elif score < 80:
            return "B"
        return "A"

users = [
    {"id": 1, "name": "Alice", "active": True, "score": 82},
    {"id": 2, "name": "Bob", "active": False, "score": 75},
    {"id": 3, "name": "Charlie", "active": True, "score": 91},
    {"id": 4, "name": "Dylan", "active": True, "score": 55},
]

example = Example()
print(example.comp(users))

# prints:
# {
# 1: {'name': 'ALICE', 'status': 'active', 'grade': 'A'}, 
# 3: {'name': 'CHARLIE', 'status': 'active', 'grade': 'A'}, 
# 4: {'name': 'DYLAN', 'status': 'active', 'grade': 'D'}
# }
Python

Set Comprehension

A Set is a unordered collection of unique elements. The syntax is almost the same as the Dictionary Comprehension, but without the key : value part.

{ expression for item in iterable }

This filters out duplicates:

from typing import List, Set
class Example:

    @staticmethod
    def comp(users: List[str]) -> Set[str]:
        return { name for name in users }


users = ["Alice", "Bob", "Alice", "Bob", "Charlie"]
print(Example.comp(users)) # prints: {'Charlie', 'Bob', 'Alice'}
Python

Filter out Alice:

from typing import List, Set
class Example:

    @staticmethod
    def comp(users: List[str]) -> Set[str]:
        return { name for name in users if name != "Alice" }


users = ["Alice", "Bob", "Alice", "Bob", "Charlie"]
print(Example.comp(users)) # prints: {'Charlie', 'Bob'}
Python

Generator Expression

This is not really the same as the other comprehensions, but for reference and relatedness I will include it here. A generator expression is similar to a list comprehension, but instead of creating the entire list in memory, it returns (yields) one item at a time. This is called a lazy evaluation and is great for large datasets.

Instead of brackets or accolades, we use parenthesis ( ):

( expression for item in iterable if condition )

Lets say we have a function that creates the Fibonacci sequence:

from typing import List
class Example:

    @staticmethod
    def fibonacci(n: int) -> List[int]:

        if n <= 0:
            raise ValueError("Must insert number great than 0")

        result = []

        for i in range(n):
            if i == 0 or i == 1:
                result.append(1)
            else:
                result.append(result[i-1] + result[i-2])

        return result


print(Example.fibonacci(10))

# Prints: [1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
Python

And lets say we want to create a generator expression out of this so we can print out each next item at a time:

f = Example.fibonacci(10)
print(f) # print [1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
generator = (n for n in f)
print(next(generator)) # print 1
print(next(generator)) # print 1
print(next(generator)) # print 2
print(next(generator)) # print 3
print(next(generator)) # print 5
print(55 in generator) # True
print(89 in generator) # False
Python

A Generator Expression creates an special type of iterator which yields items instead of returning them. Without the shortened syntax, you can write a generator with yield this way:

from typing import List
class Example:

    @staticmethod
    def count(n: int):
        max = 0
        while max < n:
            max+=1
            yield max


generator = Example.count(5)
print(next(generator))
print(next(generator))
print(next(generator))
print(next(generator))
print(next(generator))
# prints 
# 1
# 2
# 3
# 4
# 5
# next will raise an exception
Python

Iterators

Iterators have the possibility to iterate over a collection. Functions such as map(), filter(), zip() enumerate() and range() return iterators.

# MAP

def square(x):
    return x ** 2

numbers = [1, 2, 3, 4]
squares = map(square, numbers)
print(list(squares))  # prints [1, 4, 9, 16]

# using lambda is also possible
squares = map(lambda x: x**2, numbers)
print(list(squares))  # prints [1, 4, 9, 16]
Python
# FILTER

numbers = range(1, 20)
filtered = list(filter(lambda x: (x % 2 == 0) or (x % 3 == 0), numbers))
print(filtered) # prints [2, 3, 4, 6, 8, 9, 10, 12, 14, 15, 16, 18]
Python
# REDUCE

from functools import reduce

numbers = [1, 2, 3, 4]
product = reduce(lambda x, y: x * y, numbers)
print(product)  # prints 24
Python
# SORTED

numbers = [5, 1, -2, 4]
sorted_numbers = sorted(numbers, key=lambda x: abs(x))
print(sorted_numbers)  # prints [1, -2, 4, 5] (sorts by absolute value)
Python
# ZIP (Stops at shortest iterable)

first = [1, 2, 3]
second = ['a', 'b', 'c']
third = [True, False, True]

combined = zip(first, second, third)
print(list(combined))  # prints [(1, 'a', True), (2, 'b', False), (3, 'c', True)]
Python

Build-in Functions That Accept Iterators

sum()Adds all values
max() / min()Finds largest/smallest value
any() / all()Boolean checks across values
sorted()Returns a sorted list
list()Converts generator to a list
tuple()Converts generator to a tuple
enumerate()Adds index to each item
zip()Combines with other iterables

Packing and Unpacking

What I also think relates to comprehensions are packing and unpacking.

# Unpacking a Tuple

data = (10, 20, 30)
a, b, c = data
print(a, b, c)  # prints 10 20 30
Python

Functions can also use *args to pack positional arguments and **kwargs (key-worded arguments) to automatically unpack:

def greet(*names, **messages):
    for name in names:
        print(f"Hi {name}, {messages.get('msg', 'Welcome!')}")

greet("Alice", "Bob", msg="Good to see you!")
greet("Charlie", "Dylan")

# Prints:
# Hi Alice, Good to see you!
# Hi Bob, Good to see you!
# Hi Charlie, Welcome!
# Hi Dylan, Welcome!
Python

Conclusion

Python is not my most used language and I tend to forget the exact syntax of these comprehensions every now and then. With this post I hope to solidify my own knowledge and have a quick reference.

Tagged ,

Related Posts