В Python есть много возможностей, которые привлекают математиков. Вот некоторые из них: встроенная поддержка кортежей, списков и множеств, которые записываются практически так же, как это делается в математике, list comprehensions или генераторы списков, синтаксис которых похож на генераторы множеств, и другое.
Ещё один набор полезных для математиков инструментов — итераторы и генераторы, а также связанный с ними модуль 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. Мнение адмнистрации «Хекслета» может не совпадать с мнением автора оригинальной публикации.