Mastering the Itertools Module in Python

Introduction

All the tricks right on your tips

Photo by Christin Hume on Unsplash

If you have been doing python, you must have definitely come across the itertools module. It might not look like it, but I can tell you that it is one of the most powerful libraries on python. If you wish to write a more efficient, clean, and more ‘pythonic’ code, I believe that itertools is a must-have tool to add to your arsenal. Here in this article we will introduce some functionalities of the library and delve into some examples to help us better understand it.

However, the thing about itertools is that it is not enough to just know the definitions of the functions it contains. The real power lies in composing these functions to create fast, memory-efficient, and good-looking code. So here we will look at definitions along with some practical examples of the functions to understand their applications.

For an overview let us look at what the module’s official documentation has to say:

The module standardizes a core set of fast, memory efficient tools that are useful by themselves or in combination. Together, they form an “iterator algebra” making it possible to construct specialized tools succinctly and efficiently in pure Python

Itertools provides us with three different types of iterators

1. Infinite iterators
2. Terminating iterators
3. Combinatoric iterators

Let’s now dive into it

Infinite iterators

Python iterators like lists, tuples, dictionaries are exhaustive. But it is not necessary for an iterator to exhaust at some point, they can go on forever. Let’s look at the three types of infinite iterators.

1. count(start = 0, step = 1): This iterator starts returning values from ‘start’ and goes on infinitely, if steps are provided then those values are skipped. Let’s look at examples to see how it works.

   	from itertools import count
   	for i in count(0,2):
   	print(i, end = " ")
   	if i==20:
   	break
   	# OUTPUT
   	'''
   	0 2 4 6 8 10 12 14 16 18 20
   	'''

   view raw count1.py hosted with ❤ by GitHub

   We can also pass negative and floating arguments to this function.

   	from itertools import count
   	count_floating_negative = count(start=0.5, step = -0.5)
   	print(list(next(count_floating_negative) for _ in range(10)))
   	# OUTPUT
   	'''
   	[0.5, 0.0, -0.5, -1.0, -1.5, -2.0, -2.5, -3.0, -3.5, -4.0]
   	'''

   view raw count2.py hosted with ❤ by GitHub
2. cycle(iterable): It returns all values from the passed argument. When the iterable object is exhausted it starts printing from the beginning again.

   	from itertools import cycle
   	count = 0
   	for i in cycle('ABC'):
   	print(i, end = ' ')
   	count += 1
   	if count == 6:
   	break
   	# OUTPUT
   	'''
   	A B C A B C
   	'''

   view raw cycle1.py hosted with ❤ by GitHub

   Now use the next function to print the values.

3. repeat(object[, times]): This makes an iterator that returns the object over and over again. Runs indefinitely unless the time’s argument is specified.

   	from itertools import repeat
   	for i in repeat(1,5):
   	print(i, end = ' ')
   	#OUTPUT
   	'''
   	1 1 1 1 1
   	'''

   view raw repeat1.py hosted with ❤ by GitHub

   A common use of this iterator is to supply a stream of values to map() or zip().
   

   	from itertools import repeat
   	print(list(map(pow, range(10), repeat(2))))
   	#Prints the square of all numbers in range(10)
   	#OUTPUT
   	'''
   	[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

   view raw repeat2.py hosted with ❤ by GitHub

That was all about infinite iterators. From the above examples, it is clear that there may e situations where these can come in handy. Now let’s look at the real deal of the itertools library. The terminating iterators.

Terminating iterators

These iterators are used to work on finite sequences and produce an output based on the function used.

1. accumulate(iterable[, func, *, initial=None]): This makes an iterator that returns accumulated results of binary functions (specified via the optional funcargument). If func is supplied, it should be a function of two arguments. Elements of the input iterable may be any type that can be accepted as arguments to func. It will become more clear with some examples.

   	from itertools import accumulate
   	from operator import mul
   	iterable = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
   	print(list(accumulate(iterable)))
   	# Prints the addition after each iteration
   	# OUTPUT
   	"""
   	[1, 3, 6, 10, 15, 21, 28, 36, 45, 55]
   	"""
   	print(list(accumulate(iterable, mul)))
   	# Prints the mutiplication after each iteration (= factorial of the number)
   	"""
   	[1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800]
   	"""

   view raw accumulate1.py hosted with ❤ by GitHub

   Let’s pass the optional func argument to accumulate.

   	from itertools import accumulate
   	data = [-1, 2, 5, 7, -20, 9, 12, 9, 16, 4]
   	print(list(accumulate(data, min)))
   	#Prints the running minimum of the data
   	#OUTPUT
   	'''
   	[-1, -1, -1, -1, -20, -20, -20, -20, -20, -20]
   	'''
   	print(list(accumulate(data, max)))
   	#Prints the running maximum of the data
   	#OUTPUT
   	'''
   	[-1, 2, 5, 7, 7, 9, 12, 12, 16, 16]
   	'''
   	my_func = lambda a,b: (a + b) * 2
   	print(list(accumulate(data, my_func)))
   	#Prints the output calculated by the my_func
   	#OUTPUT
   	'''
   	[-1, 2, 14, 42, 44, 106, 236, 490, 1012, 2032]
   	'''

   view raw accumulate2.py hosted with ❤ by GitHub

    

   It is very similar functools.reduce() which returns only the last accumulated value, but that is a story for another day!

2. chain(*iterables): This makes an iterator that returns elements from the first iterable until it is exhausted, then proceeds to the next iterable, until all of the iterables are exhausted. In other words, it makes a single ‘sequence’ from a sequence of sequences.It has a function chain.from_iterabe() which works in a similar fashion, but it takes an argument as a list of lists instead of multiple lists.

   	from itertools import chain
   	for i in chain([0,1,2,3], range(4,11), 'ABC'):
   	print(i, end = ' ')
   	#OUTPUT
   	'''
   	0 1 2 3 4 5 6 7 8 9 10 A B C
   	'''
   	my_iterable = [[0,1,2],[3,4],[5,6,7,8,9,10],['ABC']]
   	for i in chain.from_iterable(my_iterable):
   	print(i, end = ' ')
   	#OUTPUT
   	'''
   	0 1 2 3 4 5 6 7 8 9 10 ABC
   	'''

   view raw chain.py hosted with ❤ by GitHub
3. compress(data, selectors): This makes an iterator that filters elements from data returning only those that have a corresponding element in selectors that evaluates to True.It stops when either of the two is exhausted.

   	from itertools import compress
   	data = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
   	selector = [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
   	for i in compress(data, selector):
   	print(i, end=" ")
   	# OUTPUT
   	# Simply prints all numbers for which the corresponding boolean value in selector
   	#is True, in this case all the even numbers upto 10.
   	"""
   	0 2 4 6 8 10
   	"""

   view raw compress.py hosted with ❤ by GitHub
4. dropwhile(predicate, iterable): This iterator starts returning the characters only after the predicate in argument returns false for the first time.

   	from itertools import dropwhile
   	iterable = [1, 3, 5, 7, 2, 4, 6, 9, 11]
   	predicate = lambda x: x%2==1
   	for i in dropwhile(predicate, iterable):
   	print(i, end = ' ')
   	# OUTPUT
   	"""
   	2 4 6 9 11
   	"""

   view raw dropwhile1.py hosted with ❤ by GitHub

    

5. filterfalse(predicate, iterable): This makes an iterator that filters elements from iterable returning only those for which the predicate is False. If predicate is None, it returns the items that are false.

   	from itertools import filterfalse
   	iterable = [1, 3, 5, 7, 2, 4, 6, 9, 11]
   	predicate = lambda x: x%2==0
   	for i in filterfalse(predicate, iterable):
   	print(i, end = ' ')
   	# OUTPUT
   	"""
   	1 3 5 7 9 11
   	"""

   view raw filterfalse1.py hosted with ❤ by GitHub
6. groupby(iterable, key=None): This makes an iterator that returns consecutive keys and groups from the iterable. The key is a function computing a key value for each element. If not specified or is None, key defaults to an identity function and returns the element unchanged.

   	from itertools import groupby
   	a_list = [("Even", 2),
   	("Odd", 5),
   	("Even", 8),
   	("Odd", 3)]
   	iterator = groupby(a_list, key = lambda x:x[0])
   	for key, group in iterator:
   	key_and_group = {key : list(group)}
   	print(key_and_group)
   	# OUTPUT
   	"""
   	{'Even': [('Even', 2)]}
   	{'Odd': [('Odd', 5)]}
   	{'Even': [('Even', 8)]}
   	{'Odd': [('Odd', 3)]}
   	"""

   view raw groupby1.py hosted with ❤ by GitHub

    

7. islice(iterable, start, stop[, step]): This iterator selectively returns the values from iterable container passed as the argument. Unlike regular slicing, it does not support negative values.

   	from itertools import islice
   	l = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
   	for i in islice(l, 1, 9, 1):
   	print(i, end =' ')
   	# OUTPUT
   	"""
   	2 3 4 5 6 7 8 9
   	"""

   view raw islice1.py hosted with ❤ by GitHub

    

8. starmap(function, iterable): This iterator takes a function and a list of tuples as an argument and returns the value according to the function from each tuple in the list.

   	from itertools import starmap
   	from operator import add, mul
   	iterable = [(1,3), (2,4), (3,0), (5,6)]
   	for i in starmap(add, iterable):
   	print(i, end = ' ')
   	# 4 6 3 11
   	for i in starmap(mul, iterable):
   	print(i, end = ' ')
   	# 3 8 0 30
   	for i in starmap(min, iterable):
   	print(i, end = ' ')
   	# 1 2 0 5

   view raw starmap.py hosted with ❤ by GitHub
9. takewhile(predicate, iterable): This makes an iterator that returns elements from the iterable as long as the predicate is true. It is the opposite of dropwhile.

   	from itertools import takewhile
   	iterable = [1, 3, 5, 7, 2, 4, 6, 9, 11]
   	predicate = lambda x: x%2==1
   	for i in takewhile(predicate, iterable):
   	print(i, end = ' ')
   	# OUTPUT
   	"""
   	1 3 5 7
   	"""

   view raw takewhile1.py hosted with ❤ by GitHub
10. tee(iterable, n=2): This simple returns n independent iterators from a single iterable. If n is not passed, the default value of n is 2.

    	from itertools import tee
    	iterable = [1, 2, 4, 6, 7, 13]
    	for i in tee(iterable, 4):
    	print(list(i))
    	# OUTPUT
    	"""
    	[1, 2, 4, 6, 7, 13]
    	[1, 2, 4, 6, 7, 13]
    	[1, 2, 4, 6, 7, 13]
    	[1, 2, 4, 6, 7, 13]
    	"""

    view raw tee1.py hosted with ❤ by GitHub

11. zip_longest(*iterables, fillvalue=None): This Make an iterator that aggregates elements from each of the iterables. If the iterables are of uneven length, missing values are filled-in with fillvalue. Iteration continues until the longest iterable is exhausted.

    	from itertools import zip_longest
    	iter1 = [1, 2, 3, 4, 5, 6]
    	iter2 = [1, 4, 9, 16, 25]
    	iter3 = [1, 8, 27, 64, 125]
    	for i in zip_longest(iter1, iter2, iter3, fillvalue = 0):
    	print(i)
    	# OUTPUT
    	"""
    	(1, 1, 1)
    	(2, 4, 8)
    	(3, 9, 27)
    	(4, 16, 64)
    	(5, 25, 125)
    	(6, 0, 0)
    	"""

    view raw zip_longest1.py hosted with ❤ by GitHub

These were all the terminating iterators provided by the module. You can clearly see the variety of functionalities that are provided by these iterators. These can be used in your code to make it shorter, more efficient, and also more “pythonic”.

I hope that you found these explanations easy to understand and the examples gave you the hang of it all. As you will go along, you will find more beautiful applications of these functions in solving your problems.

This wraps this article and I’ll come back with the second part of Mastering itertools where we will see the combinatoric iterators along with some more complex examples of all the functions. If you have any suggestions, feedback, or simply want to connect, I’ll be delighted to reach out. You can connect with me on LinkedIn or find some of my work on Github.

Peace.