Знакомимся с продвинутыми возможностями Python: итераторы, генераторы, itertools

В Python есть много возможностей, которые привлекают математиков. Вот некоторые из них: встроенная поддержка кортежей, списков и множеств, которые записываются практически так же, как это делается в математике, list comprehensions или генераторы списков, синтаксис которых похож на генераторы множеств, и другое.

• Итераторы
• Генераторы
• Рекурсивные генераторы
• Генераторные выражения
• Модуль itertools

Ещё один набор полезных для математиков инструментов — итераторы и генераторы, а также связанный с ними модуль itertools. Эти инструменты позволяют писать элегантный код при работе с такими математическими объектами, как бесконечные последовательности, случайные процессы, рекуррентные отношения и комбинаторные структуры. В этой публикации описана работа с итераторами и генераторами в Python.

Итераторы

Итераторы — объекты, которые позволяют обходить коллекции. Коллекции не должны обязательно существовать в памяти и быть конечными. Официальную документацию смотрите по ссылке. Примечание редактора: также смотрите наш курс по спискам.

Давайте уточним определения. Итерируемый — объект, в котором есть метод __iter__. В свою очередь, итератор — объект, в котором есть два метода: __iter__ и __next__. Почти всегда iterator возвращает себя из метода __iter__, так как они выступают итераторами для самих себя, но есть исключения.

В целом стоит избегать прямого вызова __iter__ и __next__. При использовании for или генераторов списков Python вызывает эти методы сам. Если всё-таки необходимо вызвать методы напрямую, используйте встроенные функции iter и next и в параметрах передавайте итератор или контейнер. Например, если c — итерируемый, используйте конструкцию iter(c) вместо c.__iter__(). Если a — итератор, используйте next(a), а не a.__next__(). Это похоже на использование len.

Раз уж речь зашла о len, то стоит упомянуть, что итераторы не должны иметь и часто не имеют определённой длины. Поэтому в них часто нет имплементации __len__. Чтобы подсчитать количество элементов в итераторе, приходится делать это вручную или использовать sum. Пример есть ниже после раздела о модуле itertools.

Некоторые итерируемые (iterable) не являются итераторами, но используют другие объекты как итераторы. Например, объект list относится к итерируемым, но не является итератором. В нём реализован метод __iter__, но отсутствует метод __next__. Итераторы объектов list относятся к типу listiterator. Обратите внимание, у объектов list есть определённая длина, а у listiterator нету.

>>> a = [1, 2]
>>> type(a)
<type 'list'>
>>> type(iter(a))
<type 'listiterator'>
>>> it = iter(a)
>>> next(it)
1
>>> next(it)
2
>>> next(it)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>> len(a)
2
>>> len(it)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: object of type 'listiterator' has no len()

Когда итератор завершает работу, интерпретатор Python ожидает возбуждения исключения StopIteration. Однако, как уже отмечалось, итераторы могут работать с бесконечными множествами. В таких случаях программист должен позаботиться о выходе из цикла.

Ниже дан простой пример итератора. Он считает с нуля до бесконечности. Это упрощённая версия itertools.count.

class count_iterator:
    n = 0

    def __iter__(self):
        return self

    def __next__(self):
        y = self.n
        self.n += 1
        return y

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

>>> counter = count_iterator()
>>> next(counter)
0
>>> next(counter)
1
>>> next(counter)
2
>>> next(counter)
3
>>> list(counter)  # This will result in an infinite loop!

Важная поправка к сказанному выше: если у объекта нет метода __iter__, его можно обойти, если определить метод __getitem__. В этом случае встроенная функция iter возвращает итератор с типом iterator, который использует __getitem__ для обхода элементов списка. Этот метод возвращает StopIteration или IndexError, когда обход завершается. Пример:

class SimpleList(object):
    def __init__(self, *items):
        self.items = items

    def __getitem__(self, i):
        return self.items[i]

И пример использования:

>>> a = SimpleList(1, 2, 3)
>>> it = iter(a)
>>> next(it)
1
>>> next(it)
2
>>> next(it)
3
>>> next(it)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

Рассмотрим ещё один интересный пример: генерацию последовательности Q Хофштадтера. В приведённом ниже коде итератор используется для генерации последовательности с помощью вложенных повторений.

Q(n)=Q(n−Q(n−1))+Q(n−Q(n−2))

Например, qsequence([1, 1]) генерирует точную последовательность Хофштадтера. Мы используем исключение StopIteration, чтобы показать, что последовательность не может продолжаться, так как для генерации следующего элемента должен использоваться несуществующий индекс. Если в параметрах указать значения [1, 2], последовательность немедленно заканчивается.

class qsequence:
    def __init__(self, s):
        self.s = s[:]

    def __next__(self):
        try:
            q = self.s[-self.s[-1]] + self.s[-self.s[-2]]
            self.s.append(q)
            return q
        except IndexError:
            raise StopIteration()

    def __iter__(self):
        return self

    def current_state(self):
        return self.s

Вот пример использования:

>>> Q = qsequence([1, 1])
>>> next(Q)
2
>>> next(Q)
3
>>> [next(Q) for __ in range(10)]
[3, 4, 5, 5, 6, 6, 6, 8, 8, 8]

Генераторы

Генераторами называют итераторы, определение которых выглядит как определение функций. Ещё одно определение: генераторы — функции, которые внутри используют выражение yield. Генераторы не могут возвращать значения, вместо этого выдают элементы по готовности. Python автоматизирует запоминание контекста генератора, то есть текущий поток управления, значение локальных переменных и так далее. Каждый вызов метода __next__ у объекта генератора возвращает следующее значение. Метод __iter__ также реализуется автоматически. То есть генераторы можно использовать везде, где требуются итераторы. Пример кода ниже похож на ранее рассмотренный класс итератора. Но этот код более компактный и читабельный.

def count_generator():
   n = 0
   while True:
     yield n
     n += 1

Посмотрите, как это применяется на практике.

>>> counter = count_generator()
>>> counter
<generator object count_generator at 0x106bf1aa0>
>>> next(counter)
0
>>> next(counter)
1
>>> iter(counter)
<generator object count_generator at 0x106bf1aa0>
>>> iter(counter) is counter
True
>>> type(counter)
<type 'generator'>

Теперь взгляните на реализацию последовательности Q Хофштадтера с помощью генератора. Заметьте, эта реализация значительно проще использованного выше подхода. Однако здесь уже невозможно использовать методы типа current_state. Извне невозможно получить доступ к переменным, которые хранятся в контексте генератора.

Существует словарь gi_frame.f_locals, но он относится к CPython, но не входит в стандарт языка Python.

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

def hofstadter_generator(s):
    a = s[:]
    while True:
        try:
            q = a[-a[-1]] + a[-a[-2]]
            a.append(q)
            yield q
        except IndexError:
            return

Обратите внимание, итерация в данном примере завершается простым return без параметров. Внутри происходит возбуждение исключения StopIteration. Следующий пример связан с распределением Бернулли, которое реализуется с помощью двух генераторов. Речь идёт о бесконечной последовательности случайных булевых значений. При этом вероятность True равна p, а вероятность False определяется формулой q=1-p. Затем применяется экстрактор фон Неймана, который принимает процесс Бернулли с 0 < p < 1 как источник энтропии и возвращает чистый процесс Бернулли с p = 0.5.

import random


def bernoulli_process(p):
    if p > 1.0 or p < 0.0:
        raise ValueError("p should be between 0.0 and 1.0.")

    while True:
        yield random.random() < p


def von_neumann_extractor(process):
    while True:
        x, y = next(proccess), next(process)
        if x != y:
            yield x

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

>>> def tent_map(mu, x0):
...    x = x0
...    while True:
...        yield x
...        x = mu * min(x, 1.0 - x)
...
>>>
>>> t = tent_map(2.0, 0.1)
>>> for __ in range(30):
...    print(next(t))
...
0.1
0.2
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.799999999999
0.400000000001
0.800000000003
0.399999999994
0.799999999988
0.400000000023
0.800000000047
0.399999999907
0.799999999814
0.400000000373
0.800000000745
0.39999999851
0.79999999702

Ещё один пример касается последовательности Коллатца.

def collatz(n):
   yield n
   while n != 1:
     n = n / 2 if n % 2 == 0 else 3 * n + 1
     yield n

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

Пример использования генератора:

>>> # Если гипотеза Коллатца верна, list(collatz(n)) с любым n 
... # всегда завершается
>>> list(collatz(7))
[7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]
>>> list(collatz(13))
[13, 40, 20, 10, 5, 16, 8, 4, 2, 1]
>>> list(collatz(17))
[17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]
>>> list(collatz(19))
[19, 58, 29, 88, 44, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]

Рекурсивные генераторы

Генераторы могут быть рекурсивными, как любая другая функция. Рассмотрим это на примере собственной реализации itertools.permutations. Это генератор перестановок элементов в списке. Обратите внимание, стандартная функция itertools.permutations работает быстрее, поэтому используйте её на практике. Идея простая: функция перемещает каждый элемент списка на первое место, заменяя его первым элементом в списке.

def permutations(items):
    if len(items) == 0:
        yield []
    else:
        pi = items[:]
        for i in range(len(pi)):
            pi[0], pi[i] = pi[i], pi[0]
            for p in permutations(pi[1:]):
                yield [pi[0]] + p

>>> for p in permutations([1, 2, 3]):
...     print(p)
...
[1, 2, 3]
[1, 3, 2]
[2, 1, 3]
[2, 3, 1]
[3, 1, 2]
[3, 2, 1]

Генераторные выражения

Генераторные выражения позволяют реализовать генераторы с помощью очень простого синтаксиса. В примере ниже представлен генератор, который обходит все совершенные квадраты. Результатами генераторных выражений являются объекты с типом generator. В них реализованы методы __next__ и __iter__.

>>> g = (x ** 2 for x in itertools.count(1))
>>> g
<generator object <genexpr> at 0x1029a5fa0>
>>> next(g)
1
>>> next(g)
4
>>> iter(g)
<generator object <genexpr> at 0x1029a5fa0>
>>> iter(g) is g
True
>>> [next(g) for __ in range(10)]
[9, 16, 25, 36, 49, 64, 81, 100, 121, 144]

Процесс Бернулли из предыдущих разделов также можно реализовать с помощью генераторных выражений. В данном примере p=0.4. Если в генераторном выражении требуется ещё один итератор, который используется в качестве счётчика цикла, можно использовать itertools.count для генерации бесконечной последовательности. В противном случае можно использовать range.

>>> g = (random.random() < 0.4 for __ in itertools.count())
>>> [next(g) for __ in range(10)]
[False, False, False, True, True, False, True, False, False, True]

Как отмечалось выше, генераторные выражения можно передавать в функции, которые нуждаются в итераторе. Например, сумму первых десяти совершенных квадратов можно получить так:

>>> sum(x ** 2 for x in range(10))
285

Ниже будут другие примеры генераторных выражений.

Модуль itertools

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

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

Посмотрите на интересные примеры ниже. В первом представлен альтернативный способ реализации известного паттерна: обход трёхмерного массива, а также обход всех индексов с 0≤i<j<k≤n.

from itertools import combinations, product


n = 4
d = 3


def visit(*indices):
    print(indices)

# Loop through all possible indices of a 3-D array
for i in range(n):
    for j in range(n):
        for k in range(n):
            visit(i, j, k)

# Equivalent using itertools.product
for indices in product(*([range(n)] * d)):
    visit(*indices)

# Now loop through all indices 0 <= i < j < k <= n
for i in range(n):
    for j in range(i + 1, n):
        for k in range(j + 1, n):
            visit(i, j, k)

# And equivalent using itertools.combinations
for indices in combinations(range(n), d):
    visit(*indices)

Альтернативные реализации с использованием перечислителей itertools обеспечивают два преимущества: они короче (одна строка вместо трёх), и их проще обобщать. Нужно помнить, что разница в быстродействии может быть значительной, особенно при больших значениях n.

Второй пример касается интересной математической задачи. С помощью генераторных выражений, itertools.combinations и itertools.permutations вычислим количество инверсий перестановки, а затем суммируем количество инверсий во всех перестановках в списке.

Сумма количества инверсий перестановок n равна OEIS A001809. Поэтому относительно легко показать, что закрытая форма равна n!n(n-1)/4. На практике эта формула эффективнее, но пример кода ниже необходим для отработки применения перечислителей itertools.

import itertools
import math


def inversion_number(a):
    """Return the number of inversions in list a."""
    return sum(1 for x, y in itertools.combinations(range(len(a)), 2) if a[x] > a[y])


def total_inversions(n):
    """Return total number of inversions in permutations of n."""
    return sum(inversion_number(A) for A in itertools.permutations(range(n)))

Пример использования:

>>> [total_inversions(n) for n in range(10)]
[0, 0, 1, 9, 72, 600, 5400, 52920, 564480, 6531840]

>>> [math.factorial(n) * n * (n - 1) / 4 for n in range(10)]
[0, 0, 1, 9, 72, 600, 5400, 52920, 564480, 6531840]

В качестве ещё одного примера давайте рассчитаем количество повторений с помощью метода полного перебора. Даны n и k, количество повторений Dn,k определяется как количество перестановок множества n, в котором k фиксированных значений.

Сначала пишем функцию, которая использует генераторное выражение для подсчёта фиксированных значений в перестановке. Затем используем itertools.permutations и ещё одно генераторное выражение, чтобы посчитать общее количество перестановок n, в которых k фиксированных значений. Ещё раз обратите внимание, код ниже неоптимальный, он служит для отработки применения генераторных выражений.

def count_fixed_points(p):
    """Return the number of fixed points of p as a permutation."""
    return sum(1 for x in p if p[x] == x)

def count_partial_derangements(n, k):
    """Returns the number of permutations of n with k fixed points."""
    return sum(1 for p in itertools.permutations(range(n)) if count_fixed_points(p) == k)

Применение:

# Usage:
>>> [count_partial_derangements(6, i) for i in range(7)]
[265, 264, 135, 40, 15, 0, 1]

В статье рассмотрели особенности использования итераторов, генераторов и модуля itertools в Python. Вопросы и пожелания пишите в комментариях.

Адаптированный перевод статьи A Study of Python's More Advanced Features Part I: Iterators, Generators, itertools by Sahand Saba. Мнение адмнистрации «Хекслета» может не совпадать с мнением автора оригинальной публикации.