The itertools and functools in Python

jDataLab

5 minute read

With two Python modules, itertools and functools, we can generate elegant solutions. Learn a variety of the functions from itertools and see how it would accelerate coding! functools provides higher-order functions and operations on callable objects. itertools let us compose elegant solutions for a variety of problems with the functions it provides.

pip install more-itertools
pip install functools

itertools

compress

This function filters one iterable with another.

from itertools import compress
shapes = ['circle', 'triangle', 'square', 'pentagon']
#selections = [True, False, True, False]
selections = [ 'a' in i for i in ['a','b','ac','da']  ]
result = compress(shapes, selections)#generator
list(result)

['circle', 'square', 'pentagon']

for each in result: #result can only be used one.
    print(each)

filterfalse

itertools.filterfalse(predicate, iterable): This function makes an iterator that filters elements from iterable returning only those for which the predicate is False.

from itertools import filterfalse
data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
result = filterfalse(lambda x: x<5, data)
for each in result:
    print(each)

5 6 7 8 9 10

flatten

from more_itertools import flatten
iterable = [(0, 1), (2, 3)]
list(flatten(iterable))

[0, 1, 2, 3]

chunked

from more_itertools import chunked
iterable = [0, 1, 2, 3, 4, 5, 6, 7, 8]
list(chunked(iterable, 3))

[[0, 1, 2], [3, 4, 5], [6, 7, 8]]

chain

This function takes a series of iterables and return them as one long iterable.

from itertools import chain
colors = ['red', 'orange', 'yellow', 'green', 'blue']
shapes = ['circle', 'triangle', 'square', 'pentagon']
result = chain(colors, shapes)
for each in result:
    print(each)

red orange yellow green blue circle triangle square pentagon

takewhile

itertools.takwwhile(predicate, iterable) makes an iterator and returns elements from the iterable as long as the predicate is true.

from itertools import takewhile
result = takewhile(lambda x: x<5, [1,4,6,4,1])
list(result)

[1, 4]

zip_longest

itertools.zip_longest(*iterables, fillvalue=None) makes an iterator that aggregates elements frome 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
colors = ['red', 'orange', 'yellow', 'green', 'blue',]
data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10,]
for each in zip_longest(colors, data, fillvalue=None):
    print(each)

('red', 1) ('orange', 2) ('yellow', 3) ('green', 4) ('blue', 5) (None, 6) (None, 7) (None, 8) (None, 9) (None, 10)

dropwhile

itertools.dropwhile(predicate, iterable): Make an iterator that drops elements from the iterable as long as the predicate is true; afterwards, returns every element.

from itertools import dropwhile
data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1]
result = dropwhile(lambda x: x<5, data)
for each in result:
    print(each)

5 6 7 8 9 10 1

1 < 5:  True,  drop
2 < 5:  True,  drop
3 < 5:  True,  drop
4 < 5:  True,  drop
5 < 5:  False, return surviving items

repeat

itertools.repeat(object[, times]) repeats an object over and over unless the times argument has been given.

from itertools import repeat
for i in repeat("spam", times=5):
    print(i)

spam spam spam spam spam

spy

from more_itertools import spy
iterable = (x * x for x in range(1,6))
print(iterable)
#list(iterable)

at 0x0000000055A447B0>

head, iterable = spy(iterable, n=3)
list(head)

[1, 4, 9]

list(iterable)

[1, 4, 9, 16, 25]

accumulate

from itertools import accumulate
import random
data = list(range(1,6))
print(data)

[1, 2, 3, 4, 5]

result = accumulate(data, max)
for r in result: 
  print(r)

1 2 3 4 5

random.seed(100)
data = random.sample(range(1,100), k=10)
print(data)

[19, 59, 99, 23, 91, 51, 94, 45, 56, 65]

result = accumulate(data, max)
for r in result: 
  print(r)

19 59 99 99 99 99 99 99 99 99

accumulate assumes sum if no func argument not provided.

combnations

This function takes an iterable and a integer. This will create all the unique combination that have r members.

from itertools import combinations
shapes = ['circle', 'triangle', 'square',]
result = combinations(shapes, 2)
for each in result:
  print(each)

('circle', 'triangle') ('circle', 'square') ('triangle', 'square')

combinations_with_replacement

from itertools import combinations_with_replacement
shapes = ['circle', 'triangle', 'square',]
result = combinations_with_replacement(shapes, 2)
for each in result:
    print(each)

('circle', 'circle') ('circle', 'triangle') ('circle', 'square') ('triangle', 'triangle') ('triangle', 'square') ('square', 'square')

count

Makes an iterator that returns evenly spaced values starting with number start.itertools.count(start=0, step=1)

from itertools import count
for i in count(start=10,step=3):
    print(i)
    if i > 20:
        break

10 13 16 19 22

groupby

itertools.groupby(iterable, key=None): Groups things together.

from itertools import groupby
robots = [{
    'name': 'blaster',    'faction': 1
}, {
    'name': 'jazz',    'faction': 1
},
{
    'name': 'galvatron',    'faction': 2
}, 
{
    'name': 'metroplex',    'faction': 1
}, {
    'name': 'megatron',    'faction': 2
}, {
    'name': 'starcream',    'faction': 2
}]

for key, group in groupby(robots, key=lambda x: x['faction']):
    print(key)
    print(list(group))

1 [{'name': 'blaster', 'faction': 1}, {'name': 'jazz', 'faction': 1}] 2 [{'name': 'galvatron', 'faction': 2}] 1 [{'name': 'metroplex', 'faction': 1}] 2 [{'name': 'megatron', 'faction': 2}, {'name': 'starcream', 'faction': 2}]

islice

itertools.islice(iterable, start, stop[, step]):

from itertools import islice
colors = ['red', 'orange', 'yellow', 'green', 'blue',]
few_colors = islice(colors, 2)
for each in few_colors:
    print(each)

red orange

few_colors = islice(colors, 1,3)
for each in few_colors:
    print(each)

orange yellow

permutations

from itertools import permutations
alpha_data = ['a', 'b', 'c']
result = permutations(alpha_data)
for each in result:
    print(each)

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

product

This function creates the cartesian products from a series of iterables. The set of all possible ordered pairs where the elements of A are first and the elements of B are second.

from itertools import product
num_data = [1, 2, 3]
alpha_data = ['a', 'b', 'c']
result = product(num_data, alpha_data)
for each in result:
    print(each)

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

cycle

This function cycles through an iterator endlessly.

from itertools import cycle
colors = ['red', 'orange', 'yellow', 'green', 'blue', 'violet']
for color in cycle(colors):
    print(color)