Jean's Blog

一个专注软件测试开发技术的个人博客

0%

Python应用迭代器和生成器的9个案例

发表于 分类于
学习Python学习笔记第十八讲

列表和迭代器区别

之前我们学习了列表、字典、集合等非迭代起对象，也学习了迭代器对象，觉得迭代器是多余的。

下面我们学习非迭代器和迭代器的区别。首先创建一个列表a：

1
a = [1,3,5,7]

有没有人认为，列表就是迭代器呢？可注意：列表a可不是迭代器类型（Iterator）。如果列表要想成为迭代器，需要经过内置函数iter包装：

1
a_iter = iter(a)

此时a_iter就是Iterator迭代器。可以验证：

1
2
3
4
5
6
>>> a = [1,3,5,7]
>>> a_iter = iter(a)
>>> from collections.abc import Iterator
>>> isinstance(a_iter, Iterator)
True
>>>

下面分别遍历a、a_iter：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
>>> for i in a:
...     print(i)
...
1
3
5
7
>>> for i in a_iter:
...     print(i)
...
1
3
5
7

从打印的结果开看，发现输出的结果一样，下面再次遍历a、a_iter

1
2
3
4
5
6
7
8
9
10
11
>>> for i in a:
...     print(i)
...
1
3
5
7
>>> for i in a_iter:
...     print(i)
...
>>>

从打印结果看到，再次遍历a、a_iter就会不同，a正常打印，但是a_iter没有打印出任何信息。

这就是列表a和迭代器a_iter的区别：

  • 列表不论遍历多少次，表头位置始终是第一个元素；
  • 迭代器遍历结束后，不再指向原来的表头位置，而是为最后元素的下一个位置。

只有迭代器对象才能与内置函数next结合使用，next一次，迭代器就前进一次，指向一个新的元素。

所以，要想迭代器a_iter重新指向a的表头，需要重新创建一个新的迭代a_iter_copy：

1
a_iter_copy = iter(a)

调用next，输出迭代器指向a的第一个元素：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
>>> a_iter_copy = iter(a)
>>> next(a_iter_copy)
1
>>> next(a_iter_copy)
3
>>> next(a_iter_copy)
5
>>> next(a_iter_copy)
7
>>> next(a_iter_copy)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration
>>>

值得注意，迭代器是无法用len获取长度的，只能迭代到最后一个末尾元素时，才知道其长度。那么，怎么知道迭代器到元素末尾了呢？从上面的输出结果来看，一直next，等迭代到最后一个元素后，再执行next，会触发StopIteration 异常。

所以，通过捕获此异常，就能求出迭代器指向a的长度，如下：

1
2
3
4
5
6
7
8
9
10
11
a = [1, 3, 5, 7]
a_iter = iter(a)
iter_len = 0
try:
    while True:
        i = next(a_iter)
        print(i)
        iter_len += 1
except StopIteration:
    print('iterator stops')
print('length of iterator is %d' % (iter_len,))

打印结果

1
2
3
4
5
6
7
8
9
/Users/lvjing/PycharmProjects/python_base_project/venv/bin/python /Users/lvjing/PycharmProjects/python_base_project/iter_sample_02.py
1
3
5
7
iterator stops
length of iterator is 4

Process finished with exit code 0

以上总结：

  • 遍历列表，表头位置始终不变
  • 遍历迭代器，表头位置相应改变
  • next函数执行一次，迭代对象指向就前进一次
  • StopIteration触发时，意味着已到迭代器尾部

节省内存案例

上面认识了迭代器和列表等区别后，下面来说下生成器。

带yield的函数时生成器，而生成器也是一种迭代器。所以，生成器也有以上迭代器的特点。下面通过案例来说明生成器带来哪些好处，实际的使用场景在哪里。

求一组数据累积乘，比如：三个数[1,2,3]，累积乘后返回[1,2,6]

  • 一般的方法

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    def accumulate_div(a):
        if a is None or len(a) == 0:
            return []
        rtn = [a[0]]
        for i in a[1:]:
            rtn.append(i * rtn[-1])
        return rtn


    rtn = accumulate_div([1, 2, 3, 4, 5, 6])
    print(rtn)

    执行结果

    1
    2
    3
    4
    5
    /Users/lvjing/PycharmProjects/python_base_project/venv/bin/python /Users/lvjing/PycharmProjects/python_base_project/yeild_sample_06.py
    [1]
    [1, 2, 6, 24, 120, 720]

    Process finished with exit code 0

    这个方法开辟一段新内存rtn，空间复杂度O(n)

  • 使用生成器的方法

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    def accumulate_div(a):
        if a is None or len(a) == 0:
            return []
        it = iter(a)
        total = next(it)
        yield total
        for i in it:
            total = total * i
            yield total


    # 调用生成器函数，结合for循环，输出结果
    acc = accumulate_div([1, 2, 3, 4, 5, 6])
    for i in acc:
        print(i)

    # 也可以，直接转化为list，输出结果
    rtn = list(accumulate_div([1, 2, 3, 4, 5, 6]))
    print(rtn)

    打印结果

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    /Users/lvjing/PycharmProjects/python_base_project/venv/bin/python /Users/lvjing/PycharmProjects/python_base_project/yeild_sample_07.py
    1
    2
    6
    24
    120
    720
    [1, 2, 6, 24, 120, 720]

    Process finished with exit code 0

    这种使用yield生成函数的方法，占用内存控件为O(1)。

    当输入的数组[1,2,3,4,5,6]，只有6个元素，这种内存浪费可以忽视，但是当处理几个G的数据时，这种内存控件的浪费就是致命的，尤其对于单机处理。

    学会多使用yield，写出更加节省内存的代码。

迭代案例

Python内置的itertools模块就是很好的使用yield生成器的案例，下面学习几个常用的案例。

拼接迭代器

chain函数实现元素拼接，查看官方文档中的说明（Python3.9版本），参数*表示可变的参数：

itertools.chain(*iterables)

创建一个迭代器，它首先返回第一个可迭代对象中所有元素，接着返回下一个可迭代对象中所有元素，直到耗尽所有可迭代对象中的元素。可将多个序列处理为单个序列。主要实现代码如下：

1
2
3
4
5
def chain(*iterables):
    # chain('ABC', 'DEF') --> A B C D E F
    for it in iterables:
        for element in it:
            yield element

从上面的主要代码实现，就能个看出chain是如何做到高效节省内存，chain是一个生成器函数，在迭代时，每次吐出一个元素，所以做到最高效的节省内存。下面查看实际的应用：

1
2
3
4
5
6
7
8
9
10
11
>>> from itertools import *
>>> chain_iterator = chain(['I','love'],['python'],['very','much'])
>>> for i in chain_iterator:
...     print(i)
...
I
love
python
very
much
>>>

上面的例子感觉有些像join串联字符串的感觉，但是呢，join只是一次串联一个序列对象。而chain能串联多个可迭代对象，形成一个更大的可迭代对象。

查看函数返回值chain_iterator，它是一个迭代器(Iterator)。

1
2
3
>>> from collections.abc import Iterator
>>> isinstance(chain_iterator, Iterator)
True

累积迭代器

返回可迭代对象的累积迭代器，查看官方文档中的说明（Python3.9版本）

itertools.accumulate(*iterable*[, func, *, *initial=None*])

创建一个迭代器，返回累积汇总值或其他双目运算函数的累积结果值（通过可选的 func 参数指定）。

如果提供了 func，它应当为带有两个参数的函数。 输入 iterable 的元素可以是能被 func 接受为参数的任意类型。 （例如，对于默认的加法运算，元素可以是任何可相加的类型包括 Decimal或 Fraction。）

通常，输出的元素数量与输入的可迭代对象是一致的。 但是，如果提供了关键字参数 initial，则累加会以 initial 值开始，这样输出就比输入的可迭代对象多一个元素。主要实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def accumulate(iterable, func=operator.add, *, initial=None):
    'Return running totals'
    # accumulate([1,2,3,4,5]) --> 1 3 6 10 15
    # accumulate([1,2,3,4,5], initial=100) --> 100 101 103 106 110 115
    # accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120
    it = iter(iterable)
    total = initial
    if initial is None:
        try:
            total = next(it)
        except StopIteration:
            return
    yield total
    for element in it:
        total = func(total, element)
        yield total

accumulate函数工作过程如下：

  1. 包装iterable为迭代器
  2. 初始值initial很重要
    • 如果初始值为None，迭代器向前移动求出下一个元素，并赋值给total，然后yield
    • 如果初始值被赋值，直接yield
  3. 此时迭代器it已经指向iterable的第二个元素，遍历迭代器it，func(total, element)后，求出total的第二个取值，yield后，得到返回结果的第二个元素
  4. 直到it迭代结束

下面应用的例子，返回迭代器，通过结合for打印出来。

  • 如果func不提供，默认求累积和

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    >>> accu_iterator = accumulate([1,2,3,4,5,6])
    >>> for i in accu_iterator:
    ...     print(i)
    ...
    1
    3
    6
    10
    15
    21

  • 如果func提供，func的参数个数要求是2个，根据func的累积行为返回结果

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    >>> accu_iterator = accumulate([1,2,3,4,5,6], lambda x,y: x+y)
    >>> for i in accu_iterator:
    ...     print(i)
    ...
    1
    3
    6
    10
    15
    21

漏斗迭代器

compress函数，功能类似于漏斗功能，所以称它为漏斗迭代器，查看官方文档中的说明（Python3.9版本）

itertools.compress(data, selectors)

创建一个迭代器，它返回 data 中经 selectors 真值测试为 True 的元素。迭代器在两者较短的长度处停止。主要实现代码如下：

1
2
3
def compress(data, selectors):
    # compress('ABCDEF', [1,0,1,0,1,1]) --> A C E F
    return (d for d, s in zip(data, selectors) if s)

经过selectors过滤后，返回一个更小的迭代器

1
2
3
4
5
6
7
8
>>>
>>> compress_iter = compress('abcdefg',[1,1,0,1])
>>> for i in compress_iter:
...     print(i)
...
a
b
d

compress 返回元素个数，等于两个参数中较短序列的长度。

drop迭代器

扫描可迭代对象iterable，从不满足条件处往后全部保留，返回一个更小的迭代器。查看官方文档中的说明（Python3.9版本）

itertools.dropwhile(predicate, iterable)

创建一个迭代器，如果 predicate 为true，迭代器丢弃这些元素，然后返回其他元素。注意，迭代器在 predicate 首次为false之前不会产生任何输出，所以可能需要一定长度的启动时间。主要实现代码如下：

1
2
3
4
5
6
7
8
9
def dropwhile(predicate, iterable):
    # dropwhile(lambda x: x<5, [1,4,6,4,1]) --> 6 4 1
    iterable = iter(iterable)
    for x in iterable:
        if not predicate(x):
            yield x
            break
    for x in iterable:
        yield x

函数工作过程，如下：

  • iterable包装为迭代器
  • 迭代iterable
    • 如果不满足条件predicate，yield x，然后跳出迭代，迭代完iterable剩余所有元素。
    • 如果满足条件predicate，就继续迭代，如果所有都满足，则返回空的迭代器。

应用举例：

1
2
3
4
5
6
7
8
9
10
>>> drop_iterator = dropwhile(lambda x: x<3, [1,0,2,4,1,1,3,5,-5])
>>> for i in drop_iterator:
...     print(i)
...
4
1
1
3
5
-5

take迭代器

扫描列表，只要满足条件就从可迭代对象中返回元素，直到不满足条件为止。查看官方文档中的说明（Python3.9版本）

itertools.takewhile(predicate, iterable)

创建一个迭代器，只要 predicate 为真就从可迭代对象中返回元素。主要实现代码如下：

1
2
3
4
5
6
7
def takewhile(predicate, iterable):
    # takewhile(lambda x: x<5, [1,4,6,4,1]) --> 1 4
    for x in iterable:
        if predicate(x):
            yield x
        else:
            break

函数工作过程：

  • 遍历iterable
    • 符合条件predicate, yield x
    • 否则跳出循环

应用举例：

1
2
3
4
5
6
>>> take_iterator = takewhile(lambda x: x<5, [1,4,6,4,1])
>>> for i in take_iterator:
...     print(i)
...
1
4

克隆迭代器

tee实现对原迭代器的复制。查看官方文档中的说明（Python3.9版本）

itertools.tee(iterable, n=2)

从一个可迭代对象中返回 n 个独立的迭代器。主要实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def tee(iterable, n=2):
    it = iter(iterable)
    deques = [collections.deque() for i in range(n)]
    def gen(mydeque):
        while True:
            if not mydeque:             # when the local deque is empty
                try:
                    newval = next(it)   # fetch a new value and
                except StopIteration:
                    return
                for d in deques:        # load it to all the deques
                    d.append(newval)
            yield mydeque.popleft()
    return tuple(gen(d) for d in deques)

应用案例：

克隆出2个迭代器，以元组结构返回：

1
2
3
4
5
6
7
>>> a = tee([1,4,6,4,1],2)
>>> a
(<itertools._tee object at 0x1083be6c0>, <itertools._tee object at 0x1083be700>)
>>> a[0]
<itertools._tee object at 0x1083be6c0>
>>> a[1]
<itertools._tee object at 0x1083be700>

并且，复制出两个迭代器，相互独立，如下，迭代器a[0]向前迭代一次：

1
2
>>> next(a[0])
1

克隆出的迭代器a[1]向前迭代一次，还是输出元素1，表明它们之间是相互独立的：

1
2
>>> next(a[1])
1

这种应用场景，需要用到迭代器至少两次的场合，一次迭代器用完后，再使用另一个克隆出的迭代器。

复制元素

repeat实现复制元素n次。查看官方文档中的说明（Python3.9版本）

itertools.repeat(object[, times])

创建一个迭代器，不断重复 object 。除非设定参数 times ，否则将无限重复(意思：不设置times，将会一直repeat下去。)。可用于 map() 函数中的参数，被调用函数可得到一个不变参数。也可用于 zip() 的参数以在元组记录中创建一个不变的部分。主要实现代码如下：

1
2
3
4
5
6
7
8
def repeat(object, times=None):
    # repeat(10, 3) --> 10 10 10
    if times is None:
        while True:
            yield object
    else:
        for i in range(times):
            yield object

应用如下：

1
2
3
4
>>> list(repeat(6,3))
[6, 6, 6]
>>> list(repeat([1,2,3],2))
[[1, 2, 3], [1, 2, 3]]

笛卡尔积

查看官方文档中的说明（Python3.9版本）。

itertools.product(*iterables, repeat=1)

可迭代对象输入的笛卡儿积。

大致相当于生成器表达式中的嵌套循环。例如， product(A, B) 和 ((x,y) for x in A for y in B)返回结果一样。

嵌套循环像里程表那样循环变动，每次迭代时将最右侧的元素向后迭代。这种模式形成了一种字典序，因此如果输入的可迭代对象是已排序的，笛卡尔积元组依次序发出。

要计算可迭代对象自身的笛卡尔积，将可选参数 repeat 设定为要重复的次数。例如，product(A,repeat=4) 和 product(A, A, A, A) 是一样的。

主要实现代码如下：

1
2
3
4
5
6
7
8
9
def product(*args, repeat=1):
    # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
    # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
    pools = [tuple(pool) for pool in args] * repeat
    result = [[]]
    for pool in pools:
        result = [x+[y] for x in result for y in pool]
    for prod in result:
        yield tuple(prod)

举例如下：

1
2
>>> list(product('ABCD', 'xy'))
[('A', 'x'), ('A', 'y'), ('B', 'x'), ('B', 'y'), ('C', 'x'), ('C', 'y'), ('D', 'x'), ('D', 'y')]

加强版zip

查看官方文档中的说明（Python3.9版本）。

创建一个迭代器，从每个可迭代对象中收集元素。如果可迭代对象的长度未对齐（使用repeat），将根据 fillvalue 填充缺失值。迭代持续到耗光最长的可迭代对象。主要实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def zip_longest(*args, fillvalue=None):
    # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
    iterators = [iter(it) for it in args]
    num_active = len(iterators)
    if not num_active:
        return
    while True:
        values = []
        for i, it in enumerate(iterators):
            try:
                value = next(it)
            except StopIteration:
                num_active -= 1
                if not num_active:
                    return
                iterators[i] = repeat(fillvalue)
                value = fillvalue
            values.append(value)
        yield tuple(values)

举例如下：

1
2
>>> list(zip_longest('ABCD', 'xy', fillvalue='-'))
[('A', 'x'), ('B', 'y'), ('C', '-'), ('D', '-')]