Документация Python неофициальный перевод
Содержание страницы

itertools – Функции, создающие итераторы для эффективного перебораitertools – Functions creating iterators for efficient looping


Этот модуль реализует ряд строительных блоков итераторов, вдохновлённых конструкциями из APL, Haskell и SML. Каждый из них переработан в форму, подходящую для Python.

Модуль стандартизирует основной набор быстрых и эффективных по памяти инструментов, которые полезны сами по себе или в комбинации. Вместе они образуют «алгебру итераторов», позволяющую создавать специализированные инструменты кратко и эффективно на чистом Python.

Например, SML предоставляет инструмент табуляции: tabulate(f), который порождает последовательность f(0), f(1), .... Того же эффекта можно достичь в Python скомбинировав map() и count() для получения map(f, count()).

Эти инструменты и их встроенные аналоги также хорошо работают с высокоскоростными функциями из модуля operator. Например, оператор умножения можно применить к двум векторам, чтобы получить эффективное скалярное произведение: sum(starmap(operator.mul, zip(vec1, vec2, strict=True))).

Бесконечные итераторы:

Итератор

Аргументы

Результаты

Пример

count()

start, [step]

start, start+step, start+2*step, …

count(10) --> 10 11 12 13 14 ...

cycle()

p

p0, p1, … plast, p0, p1, …

cycle('ABCD') --> A B C D A B C D ...

repeat()

elem [,n]

elem, elem, elem, … бесконечно или до n раз

repeat(10, 3) --> 10 10 10

Итераторы, завершающиеся по кратчайшей входной последовательности:

Итератор

Аргументы

Результаты

Пример

accumulate()

p [,func]

p0, p0+p1, p0+p1+p2, …

accumulate([1,2,3,4,5]) --> 1 3 6 10 15

chain()

p, q, …

p0, p1, … plast, q0, q1, …

chain('ABC', 'DEF') --> A B C D E F

chain.from_iterable()

iterable

p0, p1, … plast, q0, q1, …

chain.from_iterable(['ABC', 'DEF']) --> A B C D E F

compress()

data, selectors

(d[0] if s[0]), (d[1] if s[1]), …

compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F

dropwhile()

pred, seq

seq[n], seq[n+1], начиная с момента, когда pred не выполняется

dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1

filterfalse()

pred, seq

элементы seq, для которых pred(elem) ложно

filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8

groupby()

iterable[, key]

подитераторы, сгруппированные по значению key(v)

islice()

seq, [start,] stop [, step]

элементы из seq[start:stop:step]

islice('ABCDEFG', 2, None) --> C D E F G

pairwise()

iterable

(p[0], p[1]), (p[1], p[2])

pairwise('ABCDEFG') --> AB BC CD DE EF FG

starmap()

func, seq

func(*seq[0]), func(*seq[1]), …

starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000

takewhile()

pred, seq

seq[0], seq[1], пока pred не вернёт ложь

takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4

tee()

it, n

it1, it2, … itn разделяет один итератор на n

zip_longest()

p, q, …

(p[0], q[0]), (p[1], q[1]), …

zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-

Комбинаторные итераторы:

Итератор

Аргументы

Результаты

product()

p, q, … [repeat=1]

декартово произведение, эквивалентно вложенному циклу for

permutations()

p[, r]

кортежи длиной r, все возможные перестановки, без повторяющихся элементов

combinations()

p, r

кортежи длиной r, в отсортированном порядке, без повторяющихся элементов

combinations_with_replacement()

p, r

кортежи длиной r, в отсортированном порядке, с повторяющимися элементами

Примеры

Результаты

product('ABCD', repeat=2)

AA AB AC AD BA BB BC BD CA CB CC CD DA DB DC DD

permutations('ABCD', 2)

AB AC AD BA BC BD CA CB CD DA DB DC

combinations('ABCD', 2)

AB AC AD BC BD CD

combinations_with_replacement('ABCD', 2)

AA AB AC AD BB BC BD CC CD DD

Функции итераторных инструментовItertool functions

Следующие функции модуля создают и возвращают итераторы. Некоторые из них порождают потоки бесконечной длины, поэтому обращаться к ним следует только с помощью функций или циклов, которые обрывают поток.

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

Создаёт итератор, возвращающий накопленные суммы или накопленные результаты других бинарных функций (задаваемых через необязательный аргумент func).

Если передан аргумент func, он должен быть функцией двух аргументов. Элементы входного iterable могут быть любого типа, который может быть принят как аргумент func. (Например, при операции сложения по умолчанию элементы могут быть любого суммируемого типа, включая Decimal или Fraction.)

Обычно количество элементов вывода соответствует входному итерабельному объекту. Однако если указан именованный аргумент initial, накопление начинается со значения initial, так что в выводе оказывается на один элемент больше, чем во входном итерабельном объекте.

Примерно эквивалентно:

def accumulate(iterable, func=operator.add, *, initial=None):
    'Return running totals'
    # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
    # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115
    # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
    it = iter(iterable)
    total = initial
    if initial is None:
        try:
            total = next(it)
        except StopIteration:
            return
    yield total
    for element in it:
        total = func(total, element)
        yield total

У аргумента func есть множество применений. Его можно задать равным min() для текущего минимума, max() для текущего максимума или operator.mul() для текущего произведения. Таблицы амортизации можно построить, накапливая проценты и применяя платежи:

>>> data = [3, 4, 6, 2, 1, 9, 0, 7, 5, 8]
>>> list(accumulate(data, operator.mul))     # текущее произведение
[3, 12, 72, 144, 144, 1296, 0, 0, 0, 0]
>>> list(accumulate(data, max))              # текущий максимум
[3, 4, 6, 6, 6, 9, 9, 9, 9, 9]

# Amortize a 5% loan of 1000 with 4 annual payments of 90
>>> cashflows = [1000, -90, -90, -90, -90]
>>> list(accumulate(cashflows, lambda bal, pmt: bal*1.05 + pmt))
[1000, 960.0, 918.0, 873.9000000000001, 827.5950000000001]

См. functools.reduce() – похожая функция, возвращающая только итоговое накопленное значение.

Новое в версии 3.2.

Изменено в версии 3.3: Добавлен необязательный параметр func.

Изменено в версии 3.8: Добавлен необязательный параметр initial.

itertools.chain(*iterables)

Создаёт итератор, который возвращает элементы из первого итерируемого объекта, пока он не будет исчерпан, затем переходит к следующему, пока все итерируемые объекты не будут исчерпаны. Используется для объединения последовательностей в одну. Приблизительный эквивалент:

def chain(*iterables):
    # chain('ABC', 'DEF') --> A B C D E F
    for it in iterables:
        for element in it:
            yield element
classmethod chain.from_iterable(iterable)

Альтернативный конструктор для chain(). Получает цепочку входных данных из одного аргумента iterable, который вычисляется лениво. Примерная реализация:

def from_iterable(iterables):
    # chain.from_iterable(['ABC', 'DEF']) --> A B C D E F
    for it in iterables:
        for element in it:
            yield element
itertools.combinations(iterable, r)

Возвращает подпоследовательности элементов длины r из входного iterable.

Кортежи комбинаций выдаются в лексикографическом порядке в соответствии с порядком входного iterable. Таким образом, если входной iterable отсортирован, выходные кортежи будут создаваться в отсортированном порядке.

Элементы считаются уникальными по их позиции, а не по значению. Поэтому если входные элементы уникальны, в каждой комбинации не будет повторяющихся значений.

Примерно эквивалентно:

def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = list(range(r))
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)

Код для combinations() можно также выразить как подпоследовательность permutations() после фильтрации записей, где элементы не находятся в отсортированном порядке (согласно их позиции во входном пуле):

def combinations(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    for indices in permutations(range(n), r):
        if sorted(indices) == list(indices):
            yield tuple(pool[i] for i in indices)

Количество возвращаемых элементов равно n! / r! / (n-r)!, когда 0 <= r <= n, или нулю, когда r > n.

itertools.combinations_with_replacement(iterable, r)

Возвращает подпоследовательности элементов длины r из входного iterable, допуская повторение отдельных элементов.

Кортежи комбинаций выдаются в лексикографическом порядке в соответствии с порядком входного iterable. Таким образом, если входной iterable отсортирован, выходные кортежи будут создаваться в отсортированном порядке.

Элементы считаются уникальными по их позиции, а не по значению. Поэтому если входные элементы уникальны, сгенерированные комбинации также будут уникальны.

Примерно эквивалентно:

def combinations_with_replacement(iterable, r):
    # combinations_with_replacement('ABC', 2) --> AA AB AC BB BC CC
    pool = tuple(iterable)
    n = len(pool)
    if not n and r:
        return
    indices = [0] * r
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != n - 1:
                break
        else:
            return
        indices[i:] = [indices[i] + 1] * (r - i)
        yield tuple(pool[i] for i in indices)

Код для combinations_with_replacement() можно также выразить как подпоследовательность product() после фильтрации записей, где элементы не находятся в отсортированном порядке (согласно их позиции во входном пуле):

def combinations_with_replacement(iterable, r):
    pool = tuple(iterable)
    n = len(pool)
    for indices in product(range(n), repeat=r):
        if sorted(indices) == list(indices):
            yield tuple(pool[i] for i in indices)

Количество возвращаемых элементов равно (n+r-1)! / r! / (n-1)!, когда n > 0.

Новое в версии 3.1.

itertools.compress(data, selectors)

Создаёт итератор, который фильтрует элементы из data, возвращая только те, для которых в selectors есть соответствующий элемент, вычисляемый как True. Останавливается, когда исчерпан либо итерабельный объект data, либо selectors. Примерно эквивалентен:

def compress(data, selectors):
    # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
    return (d for d, s in zip(data, selectors) if s)

Новое в версии 3.1.

itertools.count(start=0, step=1)

Создаёт итератор, возвращающий равномерно распределённые значения, начиная с числа start. Часто используется как аргумент для map() для генерации последовательных точек данных. Также используется с zip() для добавления порядковых номеров. Примерно эквивалентен:

def count(start=0, step=1):
    # count(10) --> 10 11 12 13 14 ...
    # count(2.5, 0.5) --> 2.5 3.0 3.5 ...
    n = start
    while True:
        yield n
        n += step

При подсчёте с числами с плавающей запятой иногда можно добиться лучшей точности, заменив код на мультипликативный, например: (start + step * i for i in count()).

Изменено в версии 3.1: Добавлен аргумент step и разрешены нецелые аргументы.

itertools.cycle(iterable)

Создаёт итератор, возвращающий элементы из итерабельного объекта и сохраняющий копию каждого. Когда итерабельный объект исчерпан, возвращает элементы из сохранённой копии. Повторяется бесконечно. Примерно эквивалентен:

def cycle(iterable):
    # cycle('ABCD') --> A B C D A B C D A B C D ...
    saved = []
    for element in iterable:
        yield element
        saved.append(element)
    while saved:
        for element in saved:
              yield element

Примечание: этот элемент набора инструментов может потребовать значительного дополнительного хранилища (в зависимости от длины итерабельного объекта).

itertools.dropwhile(predicate, iterable)

Создаёт итератор, который отбрасывает элементы из итерабельного объекта, пока предикат истинен; после этого возвращает каждый элемент. Примечание: итератор не выдаёт никакого вывода, пока предикат впервые не станет ложным, поэтому время запуска может быть большим. Примерно эквивалентен:

def dropwhile(predicate, iterable):
    # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
    iterable = iter(iterable)
    for x in iterable:
        if not predicate(x):
            yield x
            break
    for x in iterable:
        yield x
itertools.filterfalse(predicate, iterable)

Создаёт итератор, который фильтрует элементы из итерабельного объекта, возвращая только те, для которых предикат ложен. Если predicate равен None, возвращает элементы, которые являются ложными. Примерно эквивалентен:

def filterfalse(predicate, iterable):
    # filterfalse(lambda x: x%2, range(10)) --> 0 2 4 6 8
    if predicate is None:
        predicate = bool
    for x in iterable:
        if not predicate(x):
            yield x
itertools.groupby(iterable, key=None)

Make 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.

Работа groupby() аналогична фильтру uniq в Unix. Она создает разрыв или новую группу каждый раз, когда изменяется значение ключевой функции (поэтому обычно необходимо отсортировать данные с помощью той же ключевой функции). Такое поведение отличается от SQL GROUP BY, которая объединяет общие элементы независимо от порядка их ввода.

Возвращаемая группа сама является итератором, который использует общий исходный итерируемый объект с groupby(). Поскольку источник общий, при продвижении объекта groupby() предыдущая группа перестает быть видимой. Поэтому, если эти данные понадобятся позже, их следует сохранить в виде списка:

groups = []
uniquekeys = []
data = sorted(data, key=keyfunc)
for k, g in groupby(data, keyfunc):
    groups.append(list(g))      # Сохранить групповой итератор как список
    uniquekeys.append(k)

groupby() примерно эквивалентно:

class groupby:
    # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
    # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D

    def __init__(self, iterable, key=None):
        if key is None:
            key = lambda x: x
        self.keyfunc = key
        self.it = iter(iterable)
        self.tgtkey = self.currkey = self.currvalue = object()

    def __iter__(self):
        return self

    def __next__(self):
        self.id = object()
        while self.currkey == self.tgtkey:
            self.currvalue = next(self.it)    # Выход при StopIteration
            self.currkey = self.keyfunc(self.currvalue)
        self.tgtkey = self.currkey
        return (self.currkey, self._grouper(self.tgtkey, self.id))

    def _grouper(self, tgtkey, id):
        while self.id is id and self.currkey == tgtkey:
            yield self.currvalue
            try:
                self.currvalue = next(self.it)
            except StopIteration:
                return
            self.currkey = self.keyfunc(self.currvalue)
itertools.islice(iterable, stop)
itertools.islice(iterable, start, stop[, step])

Создаёт итератор, возвращающий выбранные элементы из итерируемого объекта. Если start не равен нулю, то элементы из итерируемого объекта пропускаются до достижения start. Затем элементы возвращаются последовательно, если только step не установлен больше единицы, что приводит к пропуску элементов. Если stop равен None, то итерация продолжается, пока итератор не будет исчерпан (если это вообще произойдёт); в противном случае она останавливается на указанной позиции.

Если start равен None, то итерация начинается с нуля. Если step равен None, то шаг по умолчанию равен единице.

В отличие от обычной срезки, islice() не поддерживает отрицательные значения для start, stop или step. Может использоваться для извлечения связанных полей из данных, у которых внутренняя структура была уплощена (например, многострочный отчёт может выводить поле имени на каждой третьей строке).

Примерно эквивалентно:

def islice(iterable, *args):
    # islice('ABCDEFG', 2) --> A B
    # islice('ABCDEFG', 2, 4) --> C D
    # islice('ABCDEFG', 2, None) --> C D E F G
    # islice('ABCDEFG', 0, None, 2) --> A C E G
    s = slice(*args)
    start, stop, step = s.start or 0, s.stop or sys.maxsize, s.step or 1
    it = iter(range(start, stop, step))
    try:
        nexti = next(it)
    except StopIteration:
        # Потребляет *iterable* до позиции *start*.
        for i, element in zip(range(start), iterable):
            pass
        return
    try:
        for i, element in enumerate(iterable):
            if i == nexti:
                yield element
                nexti = next(it)
    except StopIteration:
        # Потребляет до *stop*.
        for i, element in zip(range(i + 1, stop), iterable):
            pass
itertools.pairwise(iterable)

Возвращает последовательные перекрывающиеся пары, взятые из входного iterable.

Количество кортежей из двух элементов в выходном итераторе будет на один меньше, чем количество входных элементов. Он будет пустым, если входной итерируемый объект содержит менее двух значений.

Примерно эквивалентно:

def pairwise(iterable):
    # pairwise('ABCDEFG') --> AB BC CD DE EF FG
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

Новое в версии 3.10.

itertools.permutations(iterable, r=None)

Возвращает последовательные перестановки длины r элементов из iterable.

Если r не указан или равен None, то r по умолчанию равен длине iterable и генерируются все возможные перестановки полной длины.

Кортежи перестановок выдаются в лексикографическом порядке в соответствии с порядком входного iterable. То есть если входной iterable отсортирован, выходные кортежи будут созданы в отсортированном порядке.

Элементы считаются уникальными на основе их позиции, а не значения. Поэтому если входные элементы уникальны, то в перестановке не будет повторяющихся значений.

Примерно эквивалентно:

def permutations(iterable, r=None):
    # permutations('ABCD', 2) --> AB AC AD BA BC BD CA CB CD DA DB DC
    # permutations(range(3)) --> 012 021 102 120 201 210
    pool = tuple(iterable)
    n = len(pool)
    r = n if r is None else r
    if r > n:
        return
    indices = list(range(n))
    cycles = list(range(n, n-r, -1))
    yield tuple(pool[i] for i in indices[:r])
    while n:
        for i in reversed(range(r)):
            cycles[i] -= 1
            if cycles[i] == 0:
                indices[i:] = indices[i+1:] + indices[i:i+1]
                cycles[i] = n - i
            else:
                j = cycles[i]
                indices[i], indices[-j] = indices[-j], indices[i]
                yield tuple(pool[i] for i in indices[:r])
                break
        else:
            return

Код для permutations() также можно выразить как подпоследовательность product(), отфильтрованную для исключения записей с повторяющимися элементами (тех, которые находятся на одной и той же позиции во входном пуле):

def permutations(iterable, r=None):
    pool = tuple(iterable)
    n = len(pool)
    r = n if r is None else r
    for indices in product(range(n), repeat=r):
        if len(set(indices)) == r:
            yield tuple(pool[i] for i in indices)

Количество возвращаемых элементов равно n! / (n-r)!, когда 0 <= r <= n, или нулю, когда r > n.

itertools.product(*iterables, repeat=1)

Декартово произведение входных итерируемых объектов.

Примерно эквивалентно вложенным циклам for в генераторном выражении. Например, product(A, B) возвращает то же, что и ((x,y) for x in A for y in B).

Вложенные циклы работают как одометр: самый правый элемент увеличивается на каждой итерации. Этот шаблон создает лексикографический порядок, так что если входные итерируемые объекты отсортированы, кортежи произведения выдаются в отсортированном порядке.

Чтобы вычислить произведение итерируемого объекта на самого себя, укажите количество повторений с помощью необязательного именованного аргумента repeat. Например, product(A, repeat=4) эквивалентно product(A, A, A, A).

Эта функция примерно эквивалентна следующему коду, за исключением того, что реальная реализация не накапливает промежуточные результаты в памяти:

def product(*args, repeat=1):
    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
    pools = [tuple(pool) for pool in args] * repeat
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)

Перед запуском product() полностью потребляет входные итерируемые объекты, сохраняя пулы значений в памяти для генерации произведений. Соответственно, она полезна только для конечных входных данных.

itertools.repeat(object[, times])

Создает итератор, который снова и снова возвращает object. Работает бесконечно, если не указан аргумент times.

Примерно эквивалентно:

def repeat(object, times=None):
    # repeat(10, 3) --> 10 10 10
    if times is None:
        while True:
            yield object
    else:
        for i in range(times):
            yield object

Обычное применение repeat – передача потока постоянных значений в map или zip:

>>> list(map(pow, range(10), repeat(2)))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
itertools.starmap(function, iterable)

Создаёт итератор, который вычисляет функцию, используя аргументы, полученные из итерируемого объекта. Используется вместо map(), когда аргументы уже сгруппированы в кортежи из одного итерируемого объекта (когда данные были «предварительно упакованы»).

Различие между map() и starmap() аналогично различию между function(a,b) и function(*c). Примерно эквивалентно:

def starmap(function, iterable):
    # starmap(pow, [(2,5), (3,2), (10,3)]) --> 32 9 1000
    for args in iterable:
        yield function(*args)
itertools.takewhile(predicate, iterable)

Создаёт итератор, возвращающий элементы из итерируемого объекта, пока предикат истинен. Приблизительно эквивалентно:

def takewhile(predicate, iterable):
    # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
    for x in iterable:
        if predicate(x):
            yield x
        else:
            break
itertools.tee(iterable, n=2)

Возвращает n независимых итераторов из одного итерируемого объекта.

Следующий код на Python помогает объяснить, что делает tee (хотя реальная реализация сложнее и использует только одну базовую очередь FIFO):

def tee(iterable, n=2):
    it = iter(iterable)
    deques = [collections.deque() for i in range(n)]
    def gen(mydeque):
        while True:
            if not mydeque:             # когда локальная deque пуста
                try:
                    newval = next(it)   # извлечь новое значение и
                except StopIteration:
                    return
                for d in deques:        # загрузить его во все двусторонние очереди
                    d.append(newval)
            yield mydeque.popleft()
    return tuple(gen(d) for d in deques)

После того как tee() создан, исходный iterable не следует использовать где-либо ещё; в противном случае iterable может продвинуться без уведомления объектов tee.

Итераторы tee не являются потокобезопасными. RuntimeError может быть возбуждён при одновременном использовании итераторов, возвращённых одним и тем же вызовом tee(), даже если исходный iterable является потокобезопасным.

Этот инструмент итерации может потребовать значительного вспомогательного хранилища (в зависимости от объёма временных данных, которые необходимо сохранить). В общем случае, если один итератор использует большую часть или все данные до того, как запустится другой итератор, быстрее использовать list() вместо tee().

itertools.zip_longest(*iterables, fillvalue=None)

Создаёт итератор, который агрегирует элементы из каждого из итерируемых объектов. Если итерируемые объекты имеют разную длину, недостающие значения заполняются с помощью fillvalue. Итерация продолжается до тех пор, пока самый длинный итерируемый объект не будет исчерпан. Приблизительно эквивалентно:

def zip_longest(*args, fillvalue=None):
    # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
    iterators = [iter(it) for it in args]
    num_active = len(iterators)
    if not num_active:
        return
    while True:
        values = []
        for i, it in enumerate(iterators):
            try:
                value = next(it)
            except StopIteration:
                num_active -= 1
                if not num_active:
                    return
                iterators[i] = repeat(fillvalue)
                value = fillvalue
            values.append(value)
        yield tuple(values)

Если один из итерируемых объектов потенциально бесконечен, то функцию zip_longest() следует обернуть чем-то, что ограничивает количество вызовов (например, islice() или takewhile()). Если не указано, fillvalue по умолчанию равен None.

Рецепты itertoolsItertools Recipes

В этом разделе представлены рецепты для создания расширенного набора инструментов с использованием существующих itertools в качестве строительных блоков.

Основная цель рецептов itertools – образовательная. Рецепты демонстрируют различные способы мышления об отдельных инструментах – например, что chain.from_iterable связано с концепцией уплощения. Рецепты также дают идеи о том, как инструменты можно комбинировать – например, как compress() и range() могут работать вместе. Рецепты также показывают шаблоны использования itertools с модулями operator и collections, а также со встроенными itertools, такими как map(), filter(), reversed() и enumerate().

Второстепенная цель рецептов – служить инкубатором. accumulate(), compress() и pairwise() itertools изначально были рецептами. В настоящее время рецепт iter_index() тестируется, чтобы проверить, оправдает ли он себя.

Практически все эти рецепты и многие-многие другие можно установить из проекта more-itertools, который находится в Python Package Index:

python -m pip install more-itertools

Многие рецепты обеспечивают такую же высокую производительность, как и базовый набор инструментов. Превосходная производительность по памяти достигается за счёт обработки элементов по одному, а не загрузки всего итерируемого объекта в память сразу. Объём кода остаётся небольшим благодаря связыванию инструментов в функциональном стиле, что помогает исключить временные переменные. Высокая скорость сохраняется за счёт предпочтения «векторизованных» строительных блоков вместо использования циклов for и генераторов, которые создают накладные расходы интерпретатора.

import collections
import math
import operator
import random

def take(n, iterable):
    "Return first n items of the iterable as a list"
    return list(islice(iterable, n))

def prepend(value, iterable):
    "Prepend a single value in front of an iterable"
    # prepend(1, [2, 3, 4]) --> 1 2 3 4
    return chain([value], iterable)

def tabulate(function, start=0):
    "Return function(0), function(1), ..."
    return map(function, count(start))

def tail(n, iterable):
    "Return an iterator over the last n items"
    # tail(3, 'ABCDEFG') --> E F G
    return iter(collections.deque(iterable, maxlen=n))

def consume(iterator, n=None):
    "Advance the iterator n-steps ahead. If n is None, consume entirely."
    # Используйте функции, потребляющие итераторы на скорости C.
    if n is None:
        # передать весь итератор в deque нулевой длины
        collections.deque(iterator, maxlen=0)
    else:
        # перейти к пустому срезу начиная с позиции n
        next(islice(iterator, n, n), None)

def nth(iterable, n, default=None):
    "Returns the nth item or a default value"
    return next(islice(iterable, n, None), default)

def all_equal(iterable):
    "Returns True if all the elements are equal to each other"
    g = groupby(iterable)
    return next(g, True) and not next(g, False)

def quantify(iterable, pred=bool):
    "Count how many times the predicate is True"
    return sum(map(pred, iterable))

def ncycles(iterable, n):
    "Returns the sequence elements n times"
    return chain.from_iterable(repeat(tuple(iterable), n))

def batched(iterable, n):
    "Batch data into tuples of length n. The last batch may be shorter."
    # batched('ABCDEFG', 3) --> ABC DEF G
    if n < 1:
        raise ValueError('n must be at least one')
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch

def grouper(iterable, n, *, incomplete='fill', fillvalue=None):
    "Collect data into non-overlapping fixed-length chunks or blocks"
    # grouper('ABCDEFG', 3, fillvalue='x') --> ABC DEF Gxx
    # grouper('ABCDEFG', 3, incomplete='strict') --> ABC DEF ValueError
    # grouper('ABCDEFG', 3, incomplete='ignore') --> ABC DEF
    args = [iter(iterable)] * n
    if incomplete == 'fill':
        return zip_longest(*args, fillvalue=fillvalue)
    if incomplete == 'strict':
        return zip(*args, strict=True)
    if incomplete == 'ignore':
        return zip(*args)
    else:
        raise ValueError('Expected fill, strict, or ignore')

def sumprod(vec1, vec2):
    "Compute a sum of products."
    return sum(starmap(operator.mul, zip(vec1, vec2, strict=True)))

def sum_of_squares(it):
    "Add up the squares of the input values."
    # sum_of_squares([10, 20, 30]) -> 1400
    return sumprod(*tee(it))

def transpose(it):
    "Swap the rows and columns of the input."
    # transpose([(1, 2, 3), (11, 22, 33)]) --> (1, 11) (2, 22) (3, 33)
    return zip(*it, strict=True)

def matmul(m1, m2):
    "Multiply two matrices."
    # matmul([(7, 5), (3, 5)], [[2, 5], [7, 9]]) --> (49, 80), (41, 60)
    n = len(m2[0])
    return batched(starmap(sumprod, product(m1, transpose(m2))), n)

def convolve(signal, kernel):
    # См.:  https://betterexplained.com/articles/intuitive-convolution/
    # convolve(data, [0.25, 0.25, 0.25, 0.25]) --> Скользящее среднее (размытие)
    # convolve(data, [1, -1]) --> Первая конечная разность (первая производная)
    # convolve(data, [1, -2, 1]) --> Вторая конечная разность (вторая производная)
    kernel = tuple(kernel)[::-1]
    n = len(kernel)
    window = collections.deque([0], maxlen=n) * n
    for x in chain(signal, repeat(0, n-1)):
        window.append(x)
        yield sumprod(kernel, window)

def polynomial_from_roots(roots):
    """Вычисляет коэффициенты многочлена по его корням.

       (x - 5) (x + 4) (x - 3)  раскрывается в:   x³ -4x² -17x + 60
    """
    # polynomial_from_roots([5, -4, 3]) --> [1, -4, -17, 60]
    expansion = [1]
    for r in roots:
        expansion = convolve(expansion, (1, -r))
    return list(expansion)

def polynomial_eval(coefficients, x):
    """Вычисляет значение многочлена в заданной точке.

    Вычисляет с лучшей численной устойчивостью, чем метод Горнера.
    """
    # Вычислить x³ -4x² -17x + 60 при x = 2.5
    # polynomial_eval([1, -4, -17, 60], x=2.5) --> 8.125
    n = len(coefficients)
    if n == 0:
        return x * 0  # Привести ноль к типу x
    powers = map(pow, repeat(x), reversed(range(n)))
    return sumprod(coefficients, powers)

def iter_index(iterable, value, start=0):
    "Return indices where a value occurs in a sequence or iterable."
    # iter_index('AABCADEAF', 'A') --> 0 1 4 7
    try:
        seq_index = iterable.index
    except AttributeError:
        # Медленный путь для общих итерируемых объектов
        it = islice(iterable, start, None)
        i = start - 1
        try:
            while True:
                yield (i := i + operator.indexOf(it, value) + 1)
        except ValueError:
            pass
    else:
        # Быстрый путь для последовательностей
        i = start - 1
        try:
            while True:
                yield (i := seq_index(value, i+1))
        except ValueError:
            pass

def sieve(n):
    "Primes less than n"
    # sieve(30) --> 2 3 5 7 11 13 17 19 23 29
    data = bytearray((0, 1)) * (n // 2)
    data[:3] = 0, 0, 0
    limit = math.isqrt(n) + 1
    for p in compress(range(limit), data):
        data[p*p : n : p+p] = bytes(len(range(p*p, n, p+p)))
    data[2] = 1
    return iter_index(data, 1) if n > 2 else iter([])

def factor(n):
    "Prime factors of n."
    # factor(99) --> 3 3 11
    for prime in sieve(math.isqrt(n) + 1):
        while True:
            quotient, remainder = divmod(n, prime)
            if remainder:
                break
            yield prime
            n = quotient
            if n == 1:
                return
    if n > 1:
        yield n

def flatten(list_of_lists):
    "Flatten one level of nesting"
    return chain.from_iterable(list_of_lists)

def repeatfunc(func, times=None, *args):
    """Повторяет вызовы func с указанными аргументами.

    Пример:  repeatfunc(random.random)
    """
    if times is None:
        return starmap(func, repeat(args))
    return starmap(func, repeat(args, times))

def triplewise(iterable):
    "Return overlapping triplets from an iterable"
    # triplewise('ABCDEFG') --> ABC BCD CDE DEF EFG
    for (a, _), (b, c) in pairwise(pairwise(iterable)):
        yield a, b, c

def sliding_window(iterable, n):
    # sliding_window('ABCDEFG', 4) --> ABCD BCDE CDEF DEFG
    it = iter(iterable)
    window = collections.deque(islice(it, n), maxlen=n)
    if len(window) == n:
        yield tuple(window)
    for x in it:
        window.append(x)
        yield tuple(window)

def roundrobin(*iterables):
    "roundrobin('ABC', 'D', 'EF') --> A D E B F C"
    # Автор рецепта – Джордж Саккис
    num_active = len(iterables)
    nexts = cycle(iter(it).__next__ for it in iterables)
    while num_active:
        try:
            for next in nexts:
                yield next()
        except StopIteration:
            # Удалить из цикла только что исчерпанный итератор.
            num_active -= 1
            nexts = cycle(islice(nexts, num_active))

def partition(pred, iterable):
    "Use a predicate to partition entries into false entries and true entries"
    # partition(is_odd, range(10)) --> 0 2 4 6 8   and  1 3 5 7 9
    t1, t2 = tee(iterable)
    return filterfalse(pred, t1), filter(pred, t2)

def before_and_after(predicate, it):
    """ Вариант takewhile(), который позволяет полностью
        доступ к остатку итератора.

        >>> it = iter('ABCdEfGhI')
        >>> all_upper, remainder = before_and_after(str.isupper, it)
        >>> ''.join(all_upper)
        'ABC'
        >>> ''.join(remainder)     # takewhile() потерял бы 'd'
        'dEfGhI'

        Обратите внимание, что первый итератор должен быть полностью
        использован до того, как второй итератор сможет
        сгенерировать корректные результаты.
    """
    it = iter(it)
    transition = []
    def true_iterator():
        for elem in it:
            if predicate(elem):
                yield elem
            else:
                transition.append(elem)
                return
    def remainder_iterator():
        yield from transition
        yield from it
    return true_iterator(), remainder_iterator()

def subslices(seq):
    "Return all contiguous non-empty subslices of a sequence"
    # subslices('ABCD') --> A AB ABC ABCD B BC BCD C CD D
    slices = starmap(slice, combinations(range(len(seq) + 1), 2))
    return map(operator.getitem, repeat(seq), slices)

def powerset(iterable):
    "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))

def unique_everseen(iterable, key=None):
    "List unique elements, preserving order. Remember all elements ever seen."
    # unique_everseen('AAAABBBCCDAABBB') --> A B C D
    # unique_everseen('ABBcCAD', str.lower) --> A B c D
    seen = set()
    if key is None:
        for element in filterfalse(seen.__contains__, iterable):
            seen.add(element)
            yield element
        # Для дедупликации с сохранением порядка,
        # более быстрое, но не ленивое решение:
        #     yield from dict.fromkeys(iterable)
    else:
        for element in iterable:
            k = key(element)
            if k not in seen:
                seen.add(k)
                yield element
        # Для случаев, когда можно вернуть последний совпадающий элемент,
        # более быстрое, но не ленивое решение:
        #      t1, t2 = tee(iterable)
        #      yield from dict(zip(map(key, t1), t2)).values()

def unique_justseen(iterable, key=None):
    "List unique elements, preserving order. Remember only the element just seen."
    # unique_justseen('AAAABBBCCDAABBB') --> A B C D A B
    # unique_justseen('ABBcCAD', str.lower) --> A B c A D
    return map(next, map(operator.itemgetter(1), groupby(iterable, key)))

def iter_except(func, exception, first=None):
    """ Повторно вызывает функцию до возникновения исключения.

    Преобразует интерфейс вызова до исключения в интерфейс итератора.
    Как builtins.iter(func, sentinel), но использует исключение вместо
    сигнального значения для завершения цикла.

    Примеры:
        iter_except(functools.partial(heappop, h), IndexError)   # итератор очереди с приоритетом
        iter_except(d.popitem, KeyError)                         # итератор словаря без блокировки
        iter_except(d.popleft, IndexError)                       # итератор дека без блокировки
        iter_except(q.get_nowait, Queue.Empty)                   # цикл по очереди производителя
        iter_except(s.pop, KeyError)                             # итератор множества без блокировки

    """
    try:
        if first is not None:
            yield first()            # Для API баз данных, требующих первоначального приведения к db.first()
        while True:
            yield func()
    except exception:
        pass

def first_true(iterable, default=False, pred=None):
    """Возвращает первое истинное значение в итерируемом объекте.

    Если истинное значение не найдено, возвращает *default*

    Если *pred* не равен None, возвращает первый элемент
    для которого pred(item) истинно.

    """
    # first_true([a,b,c], x) --> a or b or c or x
    # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
    return next(filter(pred, iterable), default)

def nth_combination(iterable, r, index):
    "Equivalent to list(combinations(iterable, r))[index]"
    pool = tuple(iterable)
    n = len(pool)
    c = math.comb(n, r)
    if index < 0:
        index += c
    if index < 0 or index >= c:
        raise IndexError
    result = []
    while r:
        c, n, r = c*r//n, n-1, r-1
        while index >= c:
            index -= c
            c, n = c*(n-r)//n, n-1
        result.append(pool[-1-n])
    return tuple(result)