Python中的迭代器和生成器

发表于 | 更新于 | 分类于 ， ， | 阅读次数：

Python中的迭代器和生成器

本博客为Python核心技术实战的学习笔记，如需获取全文，请点击链接。

迭代器

所有的容器都是可迭代（iterable）的。迭代和枚举不同，进行迭代时，并不知道被迭代对象有的容量是多少，只是每次进行迭代时，被迭代对象都返回一个元素或者迭代完成的标志（StopIteration）。

对于可迭代的对象，可以使用iter()函数得到一个迭代器，再调用next()函数便可以实现对可迭代对象的遍历。我们常使用的for循环便是将这一过程隐式化。

1
2
3
4
5
6
items = [1, 2 ,3]
it = iter(items)
try:
4print(next(it))
except StopIteration:
4print('Stop iteration')

生成器

什么是生成器？

在C++等语言中都有迭代器的概念，而对于生成器来说，很多语言中则没有对应的实现。

生成器是懒人版本的迭代器。

在迭代器中，如果我们想要对其元素进行枚举，就需要提前生成这些元素。如下述程序：

1
2
3
4
5
6
7
8
9
10
11
import os
import psutil

# 显示当前 python 程序占用的内存大小
def show_memory_info(hint):
    pid = os.getpid()
    p = psutil.Process(pid)
    
    info = p.memory_full_info()
    memory = info.uss / 1024. / 1024
    print('{} memory used: {} MB'.format(hint, memory))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def test_iterator():
    show_memory_info('initing iterator')
    list_1 = [i for i in range(100000000)]
    show_memory_info('after iterator initiated')
    print(sum(list_1))
    show_memory_info('after sum called')

def test_generator():
    show_memory_info('initing generator')
    list_2 = (i for i in range(100000000))
    show_memory_info('after generator initiated')
    print(sum(list_2))
    show_memory_info('after sum called')

%time test_iterator()
%time test_generator()

########## 输出 ##########

initing iterator memory used: 48.9765625 MB
after iterator initiated memory used: 3920.30078125 MB
4999999950000000
after sum called memory used: 3920.3046875 MB
Wall time: 17 s
initing generator memory used: 50.359375 MB
after generator initiated memory used: 50.359375 MB
4999999950000000
after sum called memory used: 50.109375 MB
Wall time: 12.5 s

在上述程序中可以发现，迭代器提前生成了所有的元素，这会消耗大量的内存。而生成器则不会提前生成所有元素。生成器只有在调用next()函数时才会生成下一个变量，同时生成器在初始化的时候不需要运行一次生成操作，耗时更短。

在Python中，我们可以使用小括号来初始化一个生成器：

1
(i for i in range(100000000))

自定义生成器

首先给出如下代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def generator(k):
    i = 1
    while True:
        yield i ** k
        i += 1

gen_1 = generator(1)
gen_3 = generator(3)
print(gen_1)
print(gen_3)

def get_sum(n):
    sum_1, sum_3 = 0, 0
    for i in range(n):
        next_1 = next(gen_1)
        next_3 = next(gen_3)
        print('next_1 = {}, next_3 = {}'.format(next_1, next_3))
        sum_1 += next_1
        sum_3 += next_3
    print(sum_1 * sum_1, sum_3)

get_sum(8)

########## 输出 ##########

<generator object generator at 0x000001E70651C4F8>
<generator object generator at 0x000001E70651C390>
next_1 = 1, next_3 = 1
next_1 = 2, next_3 = 8
next_1 = 3, next_3 = 27
next_1 = 4, next_3 = 64
next_1 = 5, next_3 = 125
next_1 = 6, next_3 = 216
next_1 = 7, next_3 = 343
next_1 = 8, next_3 = 512
1296 1296

在上述代码中，generator()函数即创建了一个生成器。在这个函数中有一个关键字yeild，可以理解为，当程序运行到这一行时会暂停，跳出程序。跳出的目的地是next()函数，yield后跟的i ** k即为next()函数的返回值。

当每次调用next()函数时，程序将会从上一次暂停的地方继续运行，并且，其中的i变量并未被清除。

与迭代器的另一个不同之处是，迭代器是一个有限的集合，而生成器可以成为一个无限集。只要调用next()函数，生成器便会执行一次操作。

下面解决一个常见的问题，给定一个列表，求一个特定的数字在这个列表中的位置。

常规做法是暴力遍历，如下所示：

1
2
3
4
5
6
7
8
9
10
11
12
13

def index_normal(L, target):
    result = []
    for i, num in enumerate(L):
        if num == target:
            result.append(i)
    return result

print(index_normal([1, 6, 2, 4, 5, 2, 8, 6, 3, 2], 2))

########## 输出 ##########

[2, 5, 9]

使用生成器的做法如下：

1
2
3
4
5
6
7
8
9
10
11

def index_generator(L, target):
    for i, num in enumerate(L):
        if num == target:
            yield i

print(list(index_generator([1, 6, 2, 4, 5, 2, 8, 6, 3, 2], 2)))

########## 输出 ##########

[2, 5, 9]

上述代码中使用list()将生成器对象转换为了列表。

另一个问题，给定两个序列，判定第一个是不是第二个的子序列。子序列指的是：一个列表的元素在第二个列表中按照顺序出现，但是不必要挨在一起。

最简单的做法是贪心算法，维持两个指针指向两个列表的最开始，对第二个列表进行遍历，当与第一个列表的指针指向的元素匹配时，第一个列表的指针后移一位，然后第二个列表的指针继续从当前位置向后遍历，当第一个列表的指针遍历完所有的元素时，返回True，否则返回False。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def is_subsequence(a, b):
4i = 0
4j = 0
4while True:
44if a[i] == b[j]:
444i += 1
44j += 1
44if j == len(b):
444break
4if i == len(a):
44# 第一个列表被遍历完
44result = True
4else:
44result = False
4return result

如果使用迭代器和生成器，则写法如下：

1
2
3
4
5
6
7
8
9
10
11
def is_subsequence(a, b):
    b = iter(b)
    return all(i in b for i in a)

print(is_subsequence([1, 3, 5], [1, 2, 3, 4, 5]))
print(is_subsequence([1, 4, 3], [1, 2, 3, 4, 5]))

########## 输出 ##########

True
False

将上述代码展开：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

def is_subsequence(a, b):
    b = iter(b)
    print(b)

    gen = (i for i in a)
    print(gen)

    for i in gen:
        print(i)

    gen = ((i in b) for i in a)
    print(gen)

    for i in gen:
        print(i)

    return all(((i in b) for i in a))

print(is_subsequence([1, 3, 5], [1, 2, 3, 4, 5]))
print(is_subsequence([1, 4, 3], [1, 2, 3, 4, 5]))

########## 输出 ##########

<list_iterator object at 0x000001E7063D0E80>
<generator object is_subsequence.<locals>.<genexpr> at 0x000001E70651C570>
1
3
5
<generator object is_subsequence.<locals>.<genexpr> at 0x000001E70651C5E8>
True
True
True
False
<list_iterator object at 0x000001E7063D0D30>
<generator object is_subsequence.<locals>.<genexpr> at 0x000001E70651C5E8>
1
4
3
<generator object is_subsequence.<locals>.<genexpr> at 0x000001E70651C570>
True
True
False
False

其中的(i in b)语句，相当于下述代码：

1
2
3
4
while True:
    val = next(b)
    if val == i:
        yield True

首先使用iter(b)将b转换为迭代器，这样在next()函数运行时，就可以保存当前运行到的b的值，下一次运行到val = next(b)时会接着上一次的指针继续往下。

1
2
3
4
5
6
7
8
9
10
11
12

b = (i for i in range(5))

print(2 in b)
print(4 in b)
print(3 in b)

########## 输出 ##########

True
True
False

上述代码中，b是一个生成器，第一次调用2 in b生成到2的时候就停止，接着调用4 in b时，生成到4时停止，此时生成器已经消耗完毕，再次调用3 in b便会返回False。

  • 容器是可迭代对象，可迭代对象调用iter()函数可以得到迭代器。对迭代器调用next()函数可以得到下一个元素。
  • 生成器是一种特殊的迭代器，使用生成器可以降低内存占用、优化程序结构、提高程序速度。
  • 对于有限的生成器来说，当遍历完成后，再次调用next()函数会触发raise StopIteration，必须使用iter()对生成器进行复位。

迭代器和生成器都是调用一次next()返回一次值，下次调用会接着上一次的状态继续执行。in和for都可以隐式调用next()函数。

参考