10 Most Useful Itertools Methods (63/100 Days of Python)

Day 63 of the “100 Days of Python” blog post series covering itertools

Python’s itertools module provides a set of tools that can be used to create efficient iterators that can be used for a variety of purposes. Itertools is a standard library that is included with Python, so there is no need to install it separately. In this tutorial, we will cover 10 of the most commonly used itertools functions and show how they can be used to create efficient iterators.

Count

The count function is used to create an iterator that returns an infinite sequence of numbers, starting from a specified value. It takes two optional arguments: start (the starting value of the sequence, which defaults to 0), and step (the increment between values, which defaults to 1):

import itertools

for i in itertools.count(1, 2):
    print(i, end=' ')
    if i >= 10:
        break

This program would print 1 3 5 7 9 .

We can also generate numbers infinitely:

import itertools

for i in itertools.count(1):
    print(i)

This code will print an infinite sequence of numbers, starting from 1.

Cycle

The cycle function is used to create an iterator that repeats a sequence infinitely. It takes a single argument, which is the sequence to be repeated:

import itertools

colors = ['red', 'green', 'blue']

for color in itertools.cycle(colors):
    print(color)

This will print the following lines in a cycle infinitely:

red
green
blue
red
green
blue
...

Repeat

The repeat function is used to create an iterator that repeats a specified value a specified number of times. It takes two arguments: value (the value to be repeated), and times (the number of times to repeat the value). If times is not specified, the value is repeated indefinitely:

import itertools

for i in itertools.repeat('hello', 3):
    print(i)

The program will print hello 3 times:

hello
hello
hello

Chain

The chain function is used to create an iterator that returns elements from multiple sequences in sequence. It takes one or more sequences as arguments, and returns an iterator that returns the elements from each sequence in turn:

import itertools

list1 = ['a', 'b', 'c']
list2 = ['d', 'e', 'f']

for item in itertools.chain(list1, list2):
    print(item)

The program will print the lists one after (“chained”) another:

a
b
c
d
e
f

Compress

The compress function is used to create an iterator that returns the elements of a sequence that correspond to a boolean mask. It takes two arguments: data (the sequence to be compressed), and selectors (a boolean mask that determines which elements to include):

import itertools

data = ['a', 'b', 'c', 'd', 'e']
selectors = [True, False, True, False, True]

for item in itertools.compress(data, selectors):
    print(item)

Which will print:

a
c
e

DropWhile

The dropwhile function is used to create an iterator that returns the elements of a sequence after a certain condition is met. It takes two arguments: predicate (a function that takes an element and returns True or False), and iterable (the sequence to be filtered):

import itertools

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

for item in itertools.dropwhile(lambda x: x < 5, data):
    print(item)

The program above will print:

5
2
4
6

GroupBy

The groupby function is used to create an iterator that groups consecutive elements of a sequence that have the same value. It takes two arguments: iterable (the sequence to be grouped), and key (a function that takes an element and returns a value used for grouping):

import itertools

data = [1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 4]

for key, group in itertools.groupby(data):
    print(key, list(group))

Which will print a nice output:

1 [1, 1]
2 [2, 2]
3 [3, 3, 3]
4 [4, 4, 4, 4]

Permutations

The permutations function is used to create an iterator that returns all possible permutations of a sequence. It takes two optional arguments: iterable (the sequence to be permuted), and r (the length of each permutation, which defaults to the length of the iterable):

import itertools

data = ['a', 'b', 'c']

for perm in itertools.permutations(data):
    print(perm)

The output would be:

('a', 'b', 'c')
('a', 'c', 'b')
('b', 'a', 'c')
('b', 'c', 'a')
('c', 'a', 'b')
('c', 'b', 'a')

Product

The product function is used to create an iterator that returns the Cartesian product of two or more sequences. It takes two or more sequences as arguments, and returns an iterator that returns tuples of all possible combinations of the elements in the sequences:

import itertools

colors = ['red', 'green', 'blue']
sizes = ['small', 'medium', 'large']

for item in itertools.product(colors, sizes):
    print(item)

Which would print all the combinations of the lists:

('red', 'small')
('red', 'medium')
('red', 'large')
('green', 'small')
('green', 'medium')
('green', 'large')
('blue', 'small')
('blue', 'medium')
('blue', 'large')

It’s important to note here that it’s possible to use the product on more than 2 lists.

Tee

The tee function is used to create multiple independent iterators from a single iterable. It takes two arguments: iterable (the sequence to be duplicated), and n (the number of independent iterators to create, which defaults to 2):

import itertools

data = ['a', 'b', 'c', 'd', 'e']

iter1, iter2 = itertools.tee(data)

print(list(iter1))    # ['a', 'b', 'c', 'd', 'e']
print(list(iter2))    # ['a', 'b', 'c', 'd', 'e']

These are just some of the most commonly used itertools functions. The itertools module provides many other functions that can be used for a variety of purposes, so be sure to check out the Python documentation for more information.

What’s next?

• If you found this story valuable, please consider clapping multiple times (this really helps a lot!)
• Hands-on Practice: Free Python Course
• Full series: 100 Days of Python
• Previous topic: Iterables
• Next topic: Glob