[Python] itertools module

accumulate(iterable[, func, *, initial=None])

This function makes an iterator that returns accumulated sums, or accumulated results of other binary functions.

The default func is operator.add function.

from itertools import accumulate
import operator

print(list(accumulate([1, 2, 3, 4])))
# [1, 3, 6, 10]

print(list(accumulate([1, 2, 3, 4], operator.mul)))
# [1, 2, 6, 24]

print(list(accumulate([1, 2, 3, 4], initial=3)))
# [3, 4, 6, 9, 13]

chain(*iterables)

This function 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.

from_iterable(iterable) gets chained inputs from a single iterable argument that is evaluated lazily.

from itertools import chain

print(list(chain([3, 2, 4,], [6, 1, 5, 9], [8, 7])))
# [3, 2, 4, 6, 1, 5, 9, 8, 7]

print(list(chain.from_iterable(['wfr', 'gjy', 'zxn'])))
# ['w', 'f', 'r', 'g', 'j', 'y', 'z', 'x', 'n']

combinations(iterable, r)

This function returns r length subsequences of elements from the input iterable.

The combination tuples are emitted in lexicographic ordering according to the order of the input iterable. So, if the input iterable is sorted, the combination tuples will be produced in sorted order.

Elements are treated as unique based on their position, not on their value. So if the input elements are unique, there will be no repeat values in each combination.

from itertools import combinations

print(list(combinations('CAB', 2)))
# [('C', 'A'), ('C', 'B'), ('A', 'B')]

combinations_with_replacement(iterable, r)

This function returns r length subsequences of elements from the input iterable allowing individual elements to be repeated more than once.

from itertools import combinations_with_replacement

print(list(combinations_with_replacement('CAB', 2)))
# [('C', 'C'), ('C', 'A'), ('C', 'B'), ('A', 'A'), ('A', 'B'), ('B', 'B')]

Compress(data, selectors)

This function makes an iterator that filters elements from data returning only those that have a corresponding element in selectors that evaluates to True.

from itertools import compress

print(list(compress('ABCDE', [1, 0, 1, 1, 0])))
# ['A', 'C', 'D']

count(start=0, step=1)

This function makes an iterator that returns evenly spaced values starting with number start.

from itertools import count

iter_count = count(5, 2)

print(next(iter_count))
# 5
print(next(iter_count))
# 7
print(next(iter_count))
# 9

cycle(iterable)

This function makes an iterator returning elements from the iterable and saving a copy of each.

from itertools import cycle

iter_cycle = cycle([1, 2, 3])
for _ in range(5):
    print(next(iter_cycle))
# 1
# 2
# 3
# 1
# 2

dropwhile(predicate, iterable)

This function makes an iterator that drops elements from the iterable as long as the predicate is True; afterward, returns every element.

from itertools import dropwhile

print(list(dropwhile(lambda x: x > 2, [5, 4, 3, 2, 1])))
# [2, 1]

print(list(dropwhile(lambda x: x > 2, [1, 2, 3, 4, 5])))
# [1, 2, 3, 4, 5]

print(list(dropwhile(lambda x: x % 2 == 0, [2, 3, 4, 5]))) 
# [3, 4, 5]

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

print(list(filterfalse(lambda x: x > 2, [1, 2, 3, 4, 5])))
# [1, 2]

print(list(filterfalse(lambda x: x % 2 == 0, [5, 4, 3, 2, 1]))) 
# [5, 3, 1]

groupby(iterable, key=None)

This function 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.

Generally, the iterable needs to already be sorted on the same key function.

from itertools import groupby

for key, group in groupby([('key1', 1), ('key1', 3), ('key2', 2), ('key2', 4)], lambda x: x[0]):
    print(f'{key}: {list(group)}')
# key1: [('key1', 1), ('key1', 3)]
# key2: [('key2', 2), ('key2', 4)]

islice(iterable, start, stop[, step])

This function makes an iterator that returns selected elements from the iterable.

If start is non-zero, then elements from the iterable are skipped until start is reached.

If stop is None, then iteration continues until the iterator is exhausted, if at all; otherwise, it stops at the specified position.

Unlike regular slicing, islice does not support negative values for start, stop, or step.

from itertools import islice

print(list(islice('ABCDE', 3)))
# ['A', 'B', 'C']

print(list(islice('ABCDE', 2, 4)))
# ['C', 'D']

print(list(islice('ABCDE', 0, None, 2)))
# ['A', 'C', 'E']

permutations(iterable, r=None)

This function returns successive r length permutations of elements in the iterable.

If r is not specified or is None, then r defaults to the length of the iterable, and all possible full-length permutations are generated.

from itertools import permutations

print(list(permutations('ABC', 2)))
# [('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'C'), ('C', 'A'), ('C', 'B')]

print(list(permutations('ABC')))
# [('A', 'B', 'C'), ('A', 'C', 'B'), ('B', 'A', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B'), ('C', 'B', 'A')]

product(*iterables, repeat=1)

This function does the Cartesian product of input iterables.

from itertools import product

print(list(product('ABC', [1, 2])))
# [('A', 1), ('A', 2), ('B', 1), ('B', 2), ('C', 1), ('C', 2)]

print(list(product([1, 2], repeat=3)))
# [(1, 1, 1), (1, 1, 2), (1, 2, 1), (1, 2, 2), (2, 1, 1), (2, 1, 2), (2, 2, 1), (2, 2, 2)]

repeat(object[, times])

This function makes an iterator that returns object over and over again.

from itertools import repeat

print(list(repeat('ABC', 3)))
# ['ABC', 'ABC', 'ABC']

starmap(function, iterable)

This function makes an iterator that computes the function using arguments obtained from the iterable.

from itertools import starmap

print(list(starmap(lambda x, y: x * y, [(1, 5), (2, 6), (3, 7), (4, 8)])))
# [5, 12, 21, 32]

takewhile(predicate, iterable)

This function makes an iterator that returns elements from the iterable as long as the predicate is True.

from itertools import takewhile

print(list(takewhile(lambda x: x > 2, [5, 4, 3, 2, 1])))
# [5, 4, 3]

print(list(takewhile(lambda x: x > 2, [1, 2, 3, 4, 5])))
# []

tee(iterable, n=2)

This function returns n independent iterators from a single iterable.

tee iterators are not threadsafe.

from itertools import tee

iter_tee = tee([1, 2, 3, 4, 5])

print(list(iter_tee[0]))
# [1, 2, 3, 4, 5]
print(list(iter_tee[1]))
# [1, 2, 3, 4, 5]

iter_tee = tee([1, 3, 5], 3)

print(list(iter_tee[0]))
# [1, 3, 5]
print(list(iter_tee[1]))
# [1, 3, 5]
print(list(iter_tee[2]))
# [1, 3, 5]

zip_longest(*iterables, fillvalue=None)

This function makes 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

print(list(zip_longest('ABCDE', [1, 2, 3])))
# [('A', 1), ('B', 2), ('C', 3), ('D', None), ('E', None)]

print(list(zip_longest('ABCDE', [1, 2, 3], fillvalue='@')))
# [('A', 1), ('B', 2), ('C', 3), ('D', '@'), ('E', '@')]