Skip to main content

迭代器 [itertools]

迭代器与生成器在内存优化上很有意义。

迭代器​

迭代器是一个可以记住遍历的位置的对象。迭代器对象从集合的第一个元素开始访问，直到所有的元素被访问完结束。迭代器只能往前不会后退。

字符串，列表或元组对象都可用于创建迭代器：

迭代器最显著的特征是拥有 __iter__() 和 __next__() 方法；它像一个链表。如果它指向末尾，那么再次执行 __next__() 时会報錯，觸發StopIteration异常，。一个例子：

a = [1, 2, 3]
b = iter(a)
print(b.__next__(), b.__next__())  # 或者使用 next(b)
1 2

实际上，Python 3 内置了一个 itertools 的库，里面有诸如 cycle 和 count 等适用于迭代器的函数：

import itertools

# count: 给定首项与公差的无穷等差数列
p = itertools.count(start = 1, step = 0.5)
print(p.__next__(), p.__next__())

# cycle: 周期循环的无穷序列
p = itertools.cycle(list("AB"))
print(next(p), next(p), next(p))

# islice: 从无穷序列中切片
p = itertools.cycle(list("AB"))
print(list(itertools.islice(p, 0, 4)))
1 1.5
A B A
['A', 'B', 'A', 'B']

请时刻注意当前指向的迭代器位置

StopIteration​

StopIteration 异常用于标识迭代的完成，防止出现无限循环的情况，在 next () 方法中我们可以设置在完成指定循环次数后触发 StopIteration 异常来结束迭代。

class MyNumbers:
  def __iter__(self):
    self.a = 1
    return self
 
  def __next__(self):
    if self.a <= 20:
      x = self.a
      self.a += 1
      return x
    else:
      raise StopIteration
 
myclass = MyNumbers()
myiter = iter(myclass)
 
for x in myiter:
  print(x)

生成器​

生成器是迭代器的一种，其实质是定义中含有 yield 关键词的函数。它没有 return() 语句。调用一个生成器函数，返回的是一个迭代器对象。

当在生成器函数中使用 yield 语句时，函数的执行将会暂停，并将 yield 后面的表达式作为当前迭代的值返回。

生成器可以直接使用类似列表解析的方式，称为生成器解析。例如：(i for i in range(10)。

def Fib(N):  # 斐波那契数列
    n, former, later = 0, 0, 1
    while n < N:
        yield later
        former, later = later, later + former
        n += 1

list(Fib(5))
[1, 1, 2, 3, 5]

上例与普通的写法看上去差别不大，但实际上可以将 while 语句改写为 while True，删除变量 n，在外部借助 itertools 的 islice 函数来截取。这在函数定义时对代码的压缩是显然的。

生成器函数 - 斐波那契

def iterFib():
    former, later = 0, 1
    while True:
        yield later
        former, later = later, later + former

list(itertools.islice(iterFib(), 0, 5))
[1, 1, 2, 3, 5]

每次调用生成器的 next() 方法或使用 for 循环进行迭代时，函数会从上次暂停的地方继续执行，直到再次遇到 yield 语句。这样，生成器函数可以逐步产生值，而不需要一次性计算并返回所有结果。

def countdown(n):
    while n > 0:
        yield n
        n -= 1
 
# 创建生成器对象
generator = countdown(5)
 
# 通过迭代生成器获取值
print(next(generator))  # 输出: 5
 
# 使用 for 循环迭代生成器
for value in generator:
    print(value)  # 输出: 4 3 2 1