Skip to content

#309: Repeating and Combining Lists With itertools

One of the things I like the most of Advent of Code is that it forces us to learn more about the language we use. This year I had to do a lot of combinations of elements in a list and for that we can either write the boring code on our own or use the itertools module that ships with Python. In this post we explore the most useful features of this hidden gem.

No installation needed

The itertools module is part of Python. Therefore, we can skip the installation and directly import it in our scripts:

1
2
3
4
5
# all, to use with itertools.cycle()
import itertools

# a specific function to use as cycle()
from itertools import cycle

Cycle through a list forever

If we need to go through our list and start over when we reach the end, we can use the cycle() function to create an endless iterator:

1
2
3
4
5
from itertools import cycle

items = [1, 2, 3, 4]
for i in cycle(items):
    print(i)

1
2
3
4
1
2
3
4
1
...

To stop the iterator, we need to use CTRL-Z.

When we have a string, we can directly use it with the itertools and do not need to convert it into a list:

1
2
3
4
5
from itertools import cycle

items = "1234"
for i in cycle(items):
    print(i)

1
2
3
4
1
2
3
4
1
...

Cycle through a list to get X values

If we know how many values we need, we can use the generated iterator of the cycle() function like this:

1
2
3
4
5
from itertools import cycle

colors = cycle(["red", "green", "blue"])
for _ in range(7):
    print(next(colors))

red
green
blue
red
green
blue
red

That way we get a lot of flexibility while keeping our code short.

Combine all elements of a single list

If we need to combine every element in a list with each other and the order does not matter, we can use the combinations() function to get the desired output:

1
2
3
4
5
from itertools import combinations

items = [1, 2, 3, 4]
for combo in combinations(items, 2):
    print(combo)

(1, 2)
(1, 3)
(1, 4)
(2, 3)
(2, 4)
(3, 4)

If we need larger groups, we can increase the length parameter:

1
2
3
items = [1, 2, 3, 4]
for combo in combinations(items, 3):
    print(combo)

(1, 2, 3)
(1, 2, 4)
(1, 3, 4)
(2, 3, 4)

Combine all elements of a list where order matters

When we need a combination of all elements in our list but (1,2) and (2,1) are not the same, we can use the permutations() function:

1
2
3
4
5
from itertools import permutations

items = [1, 2, 3, 4]
for p in permutations(items, 2):
    print(p)

(1, 2)
(1, 3)
(1, 4)
(2, 1)
(2, 3)
(2, 4)
(3, 1)
(3, 2)
(3, 4)
(4, 1)
(4, 2)
(4, 3)

Create the product of two lists

If we have two lists and want to create the Cartesian product (all tuples combining one element from each iterable), we can use the product() function:

1
2
3
4
from itertools import product

for p in product([1, 2], ["a", "b"]):
    print(p)

(1, 'a')
(1, 'b')
(2, 'a')
(2, 'b')

Conclusion

Combining values of lists or creating iterators that repeat the values for as long as we need them may not be things we need every day. But when we have a problem that requires such a behaviour, it is good to know that Python has functions for exactly that. The functions covered by this blog post are just the most useful ones from the itertools module. I strongly recommend you check for yourself what other helpful functions this module offers - it may save you a lot of time.