May 20, 2024

Python Itertools Tutorial: Efficient Data Techniques

By Alyce Osbourne

Have you ever found yourself wrestling with complex loops or struggling to efficiently manage multiple data streams in Python? If so, you’re not alone.

Well, let me show you one of my favorite modules for handling iteration!

Introducing Itertools

itertools is by far one of the most useful modules in Python’s standard library, providing a wide variety of iterators that can vastly simplify iterating over iterables. Let’s take a look at some of my most used iterators:

1. chain()

chain() simplifies the task of iterating over multiple iterables sequentially, without the need for nested loops or concatenating lists. It’s particularly handy when you’re dealing with multiple data sources that you want to process as a single sequence. This is far more efficient than concatenating iterables such as lists, as it is not a greedy operation, meaning it fits wonderfully in a generator based pipeline.

The following example combines two lists into a single iterator and prints the combined list.

import itertools
list1 = [1, 2, 3]
list2 = [4, 5, 6]
combined = itertools.chain(list1, list2)
print(list(combined))  # Output: [1, 2, 3, 4, 5, 6]

2. cycle()

cycle() is perfect for cases where you need to repeat a sequence indefinitely, such as cycling through a list of statuses or animations endlessly.

This code repeatedly cycles through a list until the loop is manually stopped after eight prints.

from itertools import cycle
count = 0
for number in cycle([1, 2, 3]):
    if count > 7:
        break
    print(number, end=' ')  # Output: 1 2 3 1 2 3 1 2
    count += 1

Note

When utilizing iterators that are essentially infinite, it is important to incorporate a break mechanism to enable a smooth exit from the loop. Neglecting to do so can cause problems with shutting down your application.

3. accumulate()

accumulate() provides a way to generate accumulated results, which can be sums, products, or any other cumulative operation. This function is useful for creating running totals, moving averages, or applying a custom accumulation function.

This will output a list of accumulated sums from the input list.

from itertools import accumulate
result = accumulate([1, 2, 3, 4])
print(list(result))  # Output: [1, 3, 6, 10]

4. combinations()

combinations() is invaluable for generating all possible combinations of elements, useful in statistics, probability, and scenarios where you need to evaluate all different groupings of items.

This example generates all possible pairs of combinations from the string “ABCD”.

from itertools import combinations
result = combinations('ABCD', 2)
print(list(result))

[('A', 'B'), ('A', 'C'), ('A', 'D'), ('B', 'C'), ('B', 'D'), ('C', 'D')]

5. permutations()

permutations() returns all possible orderings of input elements, useful for problems involving arrangements, such as scheduling or puzzle solving, as well as testing combinations of functions and classes.

The following code snippet demonstrates generating all possible permutations of the sequence ‘ABC’.

from itertools import permutations
result = permutations('ABC', 3)
print(list(result))

[('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]

6. islice()

islice() acts like slicing for lists but works on iterators, allowing you to extract specific elements without fully consuming the iterable, which is crucial for handling large data streams efficiently. Other means of slicing are generally greedy, potentially resulting in high memory overheads.

The islice function is used here to slice the range sequence similarly to list slicing.

from itertools islice
result = islice(range(10), 2, 8, 2)
print(list(result))

[2, 4, 6]

7. groupby()

groupby() is ideal for grouping adjacent items in an iterator that share a common key, which is useful for categorization or summarizing data points.

By utilizing groupby, we can simply sort the books into buckets based on the author.

from itertools import groupby

books = [
	("The Great Gatsby", "F. Scott Fitzgerald"),
    ("Tender is the Night", "F. Scott Fitzgerald"),
    ("Moby Dick", "Herman Melville"),
    ("Billy Budd", "Herman Melville"),
    ("1984", "George Orwell"),
    ("Animal Farm", "George Orwell"),
    ("Homage to Catalonia", "George Orwell"),
    ("To Kill a Mockingbird", "Harper Lee"),
    ("Go Set a Watchman", "Harper Lee")
]

books.sort(key=lambda x: x[1])

for author, group in groupby(books, key=lambda x: x[1]):
    print(author, [title for title, _ in group])

F. Scott Fitzgerald ['The Great Gatsby', 'Tender is the Night']
George Orwell ['1984', 'Animal Farm', 'Homage to Catalonia']
Harper Lee ['To Kill a Mockingbird', 'Go Set a Watchman']
Herman Melville ['Moby Dick', 'Billy Budd']

Why iterators are powerful

The power of iterators is truly unlocked when we combine them. Here are some awesome and useful combinations I have used in my projects.

1. flatten()

from itertools import chain
from typing import Iterable

def flatten[T](list_of_lists: list[list[T]]) -> Iterable[T]:
    return chain.from_iterable(list_of_lists)

This function flattens a list of lists into a single iterable.

2. ncycles()

This function will replicate and concatenate the given value N times.

from itertools import chain, repeat
from typing import Iterable

def ncycles[T](iterable: Iterable[T], n: int) -> Iterable[T]:
    return chain.from_iterable(repeat(tuple(iterable), n))

The ncycles function replicates an iterable a specified number of times.

3. dotproduct()

from itertools import starmap
from typing import Iterable
import operator

def dotproduct(vec1: list[int], vec2: list[int]) -> int:
    return sum(starmap(operator.mul, zip(vec1, vec2)))

This function calculates the dot product of two vectors using starmap and the operators module.

Note

For even more iterators, check out the more-itertools library, it massively extends the itertools module, implementing a variety of recipes to further simplify your data processing needs.

Final thoughts

These itertools functions are just the tip of the iceberg. They can simplify complex operations into more manageable, readable, and efficient solutions in Python. They can prove to be even more powerful when combined together, allowing for the iteration and aggregation of complex iterables.

For more tips, check out my guide to iterators and iterables!