Итераторы

На выражение, стоящее после for x in, питон автоматически напускает функцию iter. Она возвращает объект - итератор. Существуют и выражения-итераторы. Они выглядят как генераторы списков, но пишутся в круглых скобках, а не в квадратных. Сравним следующие 2 примера:

In [1]:
s=0
for n in [i**2 for i in range(1000)]:
    s+=n
s
Out[1]:
332833500
In [2]:
s=0
for n in (i**2 for i in range(1000)):
    s+=n
s
Out[2]:
332833500

В первом случае в памяти создаётся список из 1000 элементов. Во втором в памяти хранится только короткое выражение - итератор. Оно выдаёт очередные члены последовательности по одному, по мере надобности.

Посмотрим, как работает такое выражение.

In [3]:
it=(i**2 for i in range(4) if i!=2)
it
Out[3]:
<generator object <genexpr> at 0x7f5fa4361b48>
In [4]:
next(it)
Out[4]:
0
In [5]:
next(it)
Out[5]:
1
In [6]:
next(it)
Out[6]:
9
In [7]:
next(it)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-7-2cdb14c0d4d6> in <module>()
----> 1 next(it)

StopIteration:

Итераторы могут использоваться не только в циклах. Есть много функций с аргументами - итераторами.

In [8]:
max((10*x-x**2 for x in range(10)))
Out[8]:
25

Функция min аналогична. В таких случаях, когда выражение - итератор является единственных аргументом функции, заключать его в скобки не обязательно.

In [9]:
sum(10*x-x**2 for x in range(10))
Out[9]:
165

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

In [10]:
def f(x):
    return x**2
In [11]:
m=map(f,[0,1,2])
m
Out[11]:
<map at 0x7f5fa43953c8>
In [12]:
list(m)
Out[12]:
[0, 1, 4]

Часто бывает нужна какая-нибудь очень простая функция. Не хочется придумывать для неё имя, которое будет использовано всего 1 раз, и засорять пространство имён. В таких случаях лучше использовать анонимную функцию:

In [13]:
list(map(lambda x:2*x,[0,1,2]))
Out[13]:
[0, 2, 4]

Анонимные функции записываются так:

In [14]:
f=lambda x,y:x+2*y
f
Out[14]:
<function __main__.<lambda>>

Их, естественно, можно вызывать:

In [15]:
f(1,2)
Out[15]:
5

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

Ещё одна полезная функция - filter, она позволяет отфильтровать последовательность, оставив в ней только те элементы, которые удовлетворяют некоторому условию.

In [16]:
list(filter(lambda x:x>0,[0,1,-2,3,-4]))
Out[16]:
[1, 3]

Выражения-итераторы позволяют задавать только довольно простые последовательности. Значительно более широкие возможности предоставляют функции-генераторы. Они выглядят как функции, в которых вместо return используется yield.

In [17]:
def gen():
    yield 0
    yield -1
    yield 4

Вызвав такую функцию, мы получим некоторый объект, являющийся итератором.

In [18]:
it=gen()
it
Out[18]:
<generator object gen at 0x7f5fa42e4f68>

Его можно использовать любым обычным образом.

In [19]:
for x in it:
    print(x)

0
-1
4

Вызвав функцию gen снова, мы получим новый итератор, который опять можно использовать.

In [20]:
it=gen()
list(it)
Out[20]:
[0, -1, 4]

При первом вызове next операторы функции выполняются до первого yield. Возвращается указанное в нём значение; текущее состояние функции (точка выполнения, значения локальных переменных) запоминается. При следующем вызове next выполнение продолжается с того же места до тех пор, пока опять не встретится yield. Когда выполнение дойдёт до конца (или до return), выдаётся исключение StopIteration.

In [21]:
it=gen()
next(it)
Out[21]:
0
In [22]:
next(it)
Out[22]:
-1
In [23]:
next(it)
Out[23]:
4
In [24]:
next(it)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-24-2cdb14c0d4d6> in <module>()
----> 1 next(it)

StopIteration:

Много интересных функций для работы с итераторами имеется в модуле itertools стандартной библиотеки.

In [25]:
from itertools import repeat,count,islice,cycle,chain,accumulate

Вызов repeat(x) возвращает итератор, повторяющий значение x до бесконечности; repeat(x,n) повторяет его n раз.

In [26]:
list(repeat('abc',3))
Out[26]:
['abc', 'abc', 'abc']

Бесконечные итераторы могут использоваться для написания циклов, выход из которых производится по break; они также полезны как аргументы различных операций над итераторами. Одна из таких операций - islice: islice(it,n) - это итератор, возвращающий первые n элементов итератора it, а islice(it,n,m) возвращает элементы с n-ного (включительно) до m-го (не включая его).

In [27]:
list(islice([0,1,4,9],2))
Out[27]:
[0, 1]
In [28]:
list(islice([0,1,4,9],1,3))
Out[28]:
[1, 4]

Вызов count() возвращает итератор, выдающий бесконечную последовательность 0, 1, 2, ...; count(n) - начиная с n; count(n,h) - с шагом h.

In [29]:
list(islice(count(),10))
Out[29]:
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
In [30]:
list(islice(count(4),10))
Out[30]:
[4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
In [31]:
list(islice(count(4,2),10))
Out[31]:
[4, 6, 8, 10, 12, 14, 16, 18, 20, 22]

Вызов cycle(it) возвращает итератор, выдающий элементы it по циклу до бесконечности (для этого, разумеется, итератор it должен быть конечным; в противном случае мы никогда не доберёмся до конца первого цикла).

In [32]:
list(islice(cycle('ку'),10))
Out[32]:
['к', 'у', 'к', 'у', 'к', 'у', 'к', 'у', 'к', 'у']

Вызов chain(i1,i2) возвращает итератор, выдающий сначала все элементы i1, а затем все элементы i2. Аргументов может быть и $>2$. Разумеется, если среди аргументов встретится бесконечный итератор, то до его конца мы никогда не доберёмся.

In [33]:
for i in chain([0,1],[4,9]):
    print(i)

0
1
4
9

Есть и несколько встроенных функций для работы с итераторами, их не нужно импортировать из itertools. Так, zip работает следующим образом:

In [34]:
list(zip([0,1],[4,9]))
Out[34]:
[(0, 4), (1, 9)]

Он прекращает работу, когда закончится более короткая последовательность. Аргументов может быть и $>2$.

In [35]:
list(zip(count(),[1,2,4],'abcdefgh'))
Out[35]:
[(0, 1, 'a'), (1, 2, 'b'), (2, 4, 'c')]

Отсюда видно, что enumerate(x), который мы уже обсуждали, эквивалентен zip(count(),x).

Из числовой последовательности $x_0$, $x_1$, $x_2$, ... можно построить последовательность кумулятивных сумм $s_0=x_0$, $s_1=s_0+x_1$, $s_2=s_1+x_2$, ...

In [36]:
list(islice(accumulate(count(1)),10))
Out[36]:
[1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

Вместо сложения можно использовать любую функцию 2 переменных.

In [37]:
list(islice(accumulate(count(1),lambda x,y:x*y),10))
Out[37]:
[1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800]

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

Функция reduce из модуля functools стандартной библиотеки фактически возвращает последний элемент последовательности $s_i$, то есть $(((x_0+x_1)+x_2)+x_3)$... Вместо сложения может использоваться любая функция 2 переменных.

In [38]:
from functools import reduce
from operator import add
reduce(add,[1,4,9,16])
Out[38]:
30

С помощью итераторов можно делать поразительные вещи. Вот, например, бесконечная последовательность простых чисел (методом решета Эратосфена).

In [39]:
def primes():
    yield 2
    d={}
    for q in count(3,2):
        p=d.pop(q,None)
        if p is None: # q простое
            d[q**2]=q
            yield q
        else:         # q составное
            x=q+2*p
            while x in d:
                x+=2*p
            d[x]=p
In [40]:
list(islice(primes(),10))
Out[40]:
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29]

Можно написать класс, объекты которого являются итераторами. Сначала напишем простой класс, реализующий метод __iter__. Объект такого класса можно использовать в циклах for и других контекстах, где пребуется итератор. Там на этот объект напускается встроенная функция iter, эквивалентная вызову метода __iter__. Этот метод должен возвращать объект-итератор.

In [41]:
class List:
    
    def __init__(self,l):
        self.l=list(l)
    
    def __iter__(self):
        return ListIter(self.l)

Теперь напишем класс ListIter. Он реализует метод __next__, который вызывается, когда на объект этого класса напускается встроенная функция next. Объект хранит номер текущего лемента в списке; __next__ возвращает следующий элемент, а если его нет, возбуждает исключение StopIteration. Из одного объекта класса List можно создать сколько угодно объектов класса ListIter, которые будут пробегать по этому списку независимо друг от друга.

In [42]:
class ListIter:
    
    def __init__(self,l,n=0):
        self.l=l
        self.n=n-1
    
    def __next__(self):
        if self.n>=len(self.l)-1:
            raise StopIteration()
        else:
            self.n+=1
            return self.l[self.n]
In [43]:
l=List('xyz')
it=iter(l)
next(it)
Out[43]:
'x'
In [44]:
next(it)
Out[44]:
'y'
In [45]:
next(it)
Out[45]:
'z'
In [46]:
next(it)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-46-2cdb14c0d4d6> in <module>()
----> 1 next(it)

<ipython-input-42-a75318505da1> in __next__(self)
      7     def __next__(self):
      8         if self.n>=len(self.l)-1:
----> 9             raise StopIteration()
     10         else:
     11             self.n+=1

StopIteration:

Как мы уже обсуждали, объект класса List является итерируемым, т.е. его можно использовать в for цикле. Там не него напускается функция iter, а на получившийся объект класса ListIter - функция next, пока он не возбудит исключение StopIteration.

In [47]:
for x in l:
    print(x)

x
y
z

Можно совместить эти два класса.

In [48]:
class List2:
    
    def __init__(self,l):
        self.l=list(l)
        self.n=-1
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.n>=len(self.l)-1:
            raise StopIteration()
        else:
            self.n+=1
            return self.l[self.n]
In [49]:
l=List2('abc')
for x in l:
    print(x)

a
b
c

Но в этом случае невозможно создать несколько независимых итераторов по одному списку.

In [ ]: