Python迭代器与生成器

307 阅读5分钟

容器、可迭代对象和迭代器

容器

在Python中一切皆对象,对象的抽象就是类,而对象的集合就是容器。

列表、元组、字典、集合都是容器。对于容器,你可以很直观地想象成多个元素在一起的单元;而不同容器的区别,正是在于内部数据结构的实现方法。

可迭代对象&迭代器

所有的容器都是可迭代的(iterable)。这里的迭代,和枚举不完全一样。迭代可以想象成是你去买苹果,卖家并不告诉你他有多少库存。这样,每次你都需要告诉卖家,你要一个苹果,然后卖家采取行为:要么给你拿一个苹果;要么告诉你,苹果已经卖完了。你并不需要知道,卖家在仓库是怎么摆放苹果的。

严谨地说,迭代器(iterator)提供了一个next方法。调用这个方法后,你要么得到这个容器的下一个对象,要么得到一个StopIteration的错误。不需要像列表一样指定元素的索引,因为字典和集合这样的容器是没有索引这一说的。

而可迭代对象,通过iter()函数返回一个迭代器,再通过next()函数就可以实现遍历。for in 语句将这个过程隐式化。

如何判断一个对象是否可迭代


def is_iterable(param):
    try: 
        iter(param) 
        return True
    except TypeError:
        return False

params = [
    1234,
    '1234',
    [1, 2, 3, 4],
    set([1, 2, 3, 4]),
    {1:1, 2:2, 3:3, 4:4},
    (1, 2, 3, 4)
]
    
for param in params:
    print('{} is iterable? {}'.format(param, is_iterable(param)))

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

1234 is iterable? False
1234 is iterable? True
[1, 2, 3, 4] is iterable? True
{1, 2, 3, 4} is iterable? True
{1: 1, 2: 2, 3: 3, 4: 4} is iterable? True
(1, 2, 3, 4) is iterable? True

当然,还有另一种做法,是isinstance(obj, Iterable)。

生成器

关于生成器,你只需要记住一点:生成器是懒人版本的迭代器

我们知道,在迭代器中,如果我们想要枚举它的元素,这些元素需要事先生成。而生成器不需要,如下:


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))
    

def test_iterator():
    show_memory_info('initing iterator')
    # 声明一个迭代器,[i for i in range(100000000)]生成包含一亿元素的列表。
    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')

test_iterator()
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
initing generator memory used: 50.359375 MB
after generator initiated memory used: 50.359375 MB
4999999950000000
after sum called memory used: 50.109375 MB

由上面的代码可以看出,[i for i in range(100000000)]生成包含一亿元素的列表。每个元素在生成后都会保存到内存中,它们占有了巨量的内存,内存不够的话就会出现OOM错误。

但是我们并不需要在内存里同时保存这么多东西,有些东西用完就可以扔掉。于是,生成器出现了,在你调用next()函数的时候,才会生成下一个变量。

你可以清晰地看到,生成器并不会像迭代器一样占用大量内存,只有在被使用的时候才会调用。而且生成器在初始化的时候,并不需要运行一次性生成操作。

生成器的应用

使用生成器验证此公式:(1 + 2 + 3 + ... + n)^2 = 1^3 + 2^3 + 3^3 + ... + n^3的正确性

#返回一个生成器
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()这个函数,它返回了一个生成器。接下来的yield是关键,你可以理解为,函数运行到这一行的时候,程序会从这里暂停,然后跳出到next()函数。i**k其实成了next()函数的返回值。

这样,每次next(gen)函数被调用的时候,暂停的程序就又复活了,从yield这里向下继续执行;同时注意,局部变量i并没有被清除掉,而是会继续累加。

这个时候,你应该注意到了,这个生成器居然可以一直进行下去!没错,事实上,迭代器是一个有限集合,生成器则可以成为一个无限集。你只管调用next(),生成器根据运算会自动生成新的元素,然后返回给你。

给定一个 list 和一个指定数字,求这个数字在 list 中的位置。

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]

需要强调的是,index_generator会返回一个Generator对象,需要使用list转换为列表后,才能用print输出。

判断某些元素是否按顺序出现在可迭代对象中

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

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

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

True
True
False

next()函数运行的时候,保存了当前的指针。而指针是不会往回走的,只能是不断往前。