作者:老齐

如果你没有听说过被称为“宝石”的Python 3标准库itertools,未免是一种遗憾。

在Python的官方文档中,有关于itertools的丰富的内容,所以,从官方文档开始,是了解它的最佳途径。

对于itertools而言,不仅要知道它包含多少个函数,更要理解这些函数能够让代码更快、更省内存,也更优雅。

本文要介绍几个例子,不仅仅是函数的应用,还是通过实践操作,能够理解编程中“迭代式思维”的特点。通常,所举的例子都是由浅入深的。

请读者注意:阅读本文之前,你应该已经对Python中的迭代器和生成器有了初步认识。如果没有,阅读本文会有一些困难。所以,请先阅读**《跟老齐学Python:轻松入门》**的相关章节。

一、理解zip函数

看下面的操作过程,其中使用`zip`函数——内置函数。

>>> list(zip([1, 2, 3], ['a', 'b', 'c']))
[(1, 'a'), (2, 'b'), (3, 'c')]

[1, 2, 3] 和 ['a', 'b', 'c'] 是列表,所有列表都是可迭代的。这意味着可以一次返回一个元素。

zip函数最终得到一个zip对象——一个迭代器对象。而这个迭代器对象是由若干个元组组成的。

那么这些元组是如何生成的呢?

对照上面的代码,从开始算起。按照Python中的技术习惯,开始的那个元组是用0来计数的,即第0个。

1. 第0个元组中索引为0的元素,来自于zip函数第0个参数(`[1,2,3]`)当前指针所指的值,刚开始读取,那就是1;

2. 第0个元组中索引为1的元素,来自于zip函数中第1个参数(`['a', 'b', 'c']`)当前指针所指的值,也是刚刚开始读取,应该是a;于是组成了zip对象的第0个元组`(1, 'a')`。

3. 第0个元组中索引为2的元素,来自于zip函数中第2个参数——没有第二个,上面的参数总共有0、1两个。那么这个元组就组建完毕了,也就没有索引为2的元素。接下来再组建下一个元组,按照这里的计数顺序应该是第1个元组。

4. zip对象第1个元组中索引为0的元素,来自于zip函数第0个参数当前指针所指的值。注意,因为上次已经度去过1了,这时候对应的值是2。

5. 同上面的道理,这时候指针所指的值是'b'。于是组成了zip对象第1个元组`(2, 'b')`。

6. 如此重复,就得到了最终的结果。

请注意上面的叙述,如果把元组中的组成对象来源概括一句话,那就是元组中的第i个元素就来自于第i个参数中指针在当前所指的元素对象——请细细品味这句话的含义,后面有大用途。

对于zip函数,上面的过程,貌似“压缩”一样,那么,是否有反过程——解压缩。例如从上面示例的结果中分别恢复出来原来的两个对象。

有。一般认为是这这么做:

>>> result = zip([1, 2, 3], ["a", "b", "c"])
>>> c, v = zip(*result)
>>> print(c, v)
(1, 2, 3) ('a', 'b', 'c')

这是什么原理。

result应用的是一个迭代器对象,不过为了能够显示的明白,也可以理解为是[(1, 'a'), (2, 'b'), (3, 'c')]。接下来使用了zip(*result),这里的符号`*`的作用是收集参数,在函数的参数中,有对此详细阐述(请参阅《跟老齐学Python:轻松入门》)。

>>> def foo(*a): print(a)
...
>>> lst = [(1,2), (3,4)]
>>> foo(*lst)
((1, 2), (3, 4))
>>> foo((1,2), (3,4))
((1, 2), (3, 4))

仿照这个示例,就能明晰下面两个操作是等效的。

>>> lst = [(1, 'a'), (2, 'b'), (3, 'c')]
>>> zip(*lst)
<zip object at 0x104c8fb48>
>>> zip((1, 'a'), (2, 'b'), (3, 'c'))
<zip object at 0x104f27308>

从返回的对象内存编码也可以看出,两个是同样的对象。

既然如此,我们就可以通过理解`zip((1, 'a'), (2, 'b'), (3, 'c'))`的结果产生过程来理解`zip(*lst)`了。而前者生成结果的过程前面已经阐述过了,此处不再赘述。

原来,所谓的“解压缩”和“压缩”,计算的方法是一样的。豁然开朗。

其它常用函数

除了`zip`函数,还有一个内置函数`iter`,它以可迭代对象为参数,会返回迭代器对象。

>>> iter([1, 2, 3, 4])
<list_iterator object at 0x7fa80af0d898>

本质上,`iter`函数调用参数的每个元素,然后借助于`__next__`函数返回迭代器对象,并把结果集成到一个元组中。

内置函数`map`是另外一个返回迭代器对象的函数,它以只有一个参数的函数对象为参数,这个函数每次从可迭代对象中取一个元素。

>>> list(map(len, ['abc', 'de', 'fghi']))
[3, 2, 4]

`map`函数的执行原理是:用`__iter__`函数调用第二个参数,并用`__next__`函数返回执行结果。在上面的例子中,`len`函数要调用后面的列表中的每个元素,并返回一个迭代器对象。

既然迭代器是可迭代的,就可以把`zip`返回的迭代器对象用到`map`函数的参数中了。例如,用下面的方式计算两个列表中对应元素的和。

>>> list(map(sum, zip([1, 2, 3], [4, 5, 6])))
[5, 7, 9]

迭代器对象有两个主要功效,一是节省内存,而是提高执行效率。

典型问题

有一个列表,由一些正整数组成,如`[1, 2, 3, 4, 5, 6]`,写一个函数,函数的一个参数`n`表示要将列表中几个元素划为一组,假设n=2,则将两个元素为一组,最终返回`[(1, 2), (3, 4), (5, 6)]`。

如果用简单的方式,可以这样写此函数:

def naive_grouper(inputs, n):
    num_groups = len(inputs) // n
    return [tuple(inputs[i*n:(i+1)*n]) for i in range(num_groups)]

测试一下,此函数能够如我们所愿那样工作。

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> naive_grouper(nums, 2)
[(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]

但是,在上面的测试中,所传入的列表元素个数是比较小的,如果列表元素个数很多,比如有100万个。这就需要有比较大的内存了,否则无法执行运算。

但是,如果这样执行此程序:

def naive_grouper(inputs, n):
    num_groups = len(inputs) // n
    return [tuple(inputs[i*n:(i+1)*n]) for i in range(num_groups)]
 
for _ in naive_grouper(range(100000000), 10):
    pass

把上面的程序保存为文件`naive.py`。可以用下面的指令,测量运行程序时所占用的内存空间和耗费的时长。注意,要确保你本地机器的内存至少5G。

$ time -f "Memory used (kB): %M\nUser time (seconds): %U" python3 naive.py
Memory used (kB): 4551872
User time (seconds): 11.04

注意:Ubuntu系统中,你可能要执行 `/usr/bin/time`。

把列表或者元素传入`naïve_grouper`函数,需要计算机提供4.5GB的内存空间,才能执行`range(100000000)`的循环。

如果采用迭代器对象,就会有很大变化了。

def better_grouper(inputs, n):
    iters = [iter(inputs)] * n
    return zip(*iters)

这个简短的函数中,内涵还是很丰富的。所以我们要逐行解释。

表达式`[iters(inputs)] * n`创建一个迭代器对象,它包含了n个同样的列表对象。

下面以`n=2`为例说明。

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> iters = [iter(nums)] * 2
>>> list(id(itr) for itr in iters)  # 内存地址是一样的
[139949748267160, 139949748267160]

在`iters`中的两个迭代器对象是同一个对象——认识到这一点非常重要。

结合前面对`zip`的理解,`zip(*iters)`和`zip(iter(nums), iter(nums))`是一样的。为了能够以更直观的方式进行说明,就可以认为是`zip((1, 2, 3, 4, 5, 6, 7, 8, 9, 10), (1, 2, 3, 4, 5, 6, 7, 8, 9, 10))`。按照前文所述的`zip`工作流程,其计算过程如下:

1. 第0个元组中的第0个元素,来自于第0个参数中指针当前所指对象,是1;

2. 第0个元组中的第1个元素,来自于第1个参数中指针当前所指对象,特别注意,两个参数是同一个对象,此时指针所指的是2。

3. 第0个元组中的第2个元素,来自于第2个参数——没有。于是乎得到了元组`(1, 2)`。

4. 接下来是第1个元组中的第0个元素,来自第0个参数中指针当前所指的对象,还是因为同一个对象的元素,此时指针所指的是3。

5. 依次类推,得到(3,4)等元组。

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> list(better_grouper(nums, 2))
[(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]

上面的函数`better_grouper()`的优点在于:

- 不使用内置函数`len()`,能够以任何可迭代对象为参数
- 返回的是可迭代对象,不是列表。这样更少占用内存

把上述流程保存为文件`better.py`。

def better_grouper(inputs, n):
    iters = [iter(inputs)] * n
    return zip(*iters)
 
for _ in better_grouper(range(100000000), 10):
    pass

然后使用 `time`在终端执行。

$ time -f "Memory used (kB): %M\nUser time (seconds): %U" python3 better.py
Memory used (kB): 7224
User time (seconds): 2.48

对比前面执行 `naive.py`,不论在内存还是执行时间上,都表现非常优秀。

进一步研究

对于上面的`better_grouper`函数,深入分析一下,发现它还有问题。它只能分割能够被列表长度整除的列表中的数字,如果不能整除的话,就会有一些元素被舍弃。例如:

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> list(better_grouper(nums, 4))
[(1, 2, 3, 4), (5, 6, 7, 8)]
 
如果要4个元素一组,就会有9和10不能分组了。之所以,还是因为`zip`函数。
>>> list(zip([1, 2, 3], ['a', 'b', 'c', 'd']))
[(1, 'a'), (2, 'b'), (3, 'c')]

最终得到的zip对象中的元组数量,是参数中长度最小的对象决定。

如果你感觉这样做不爽,可以使用`itertools.zip_longest()`,这个函数是以最长的参数为基准,如果有不足的,默认用`None`填充,当然也可以通过参数`fillvalue`指定填充对象。

>>> import itertools
>>> x = [1, 2, 3]
>>> y = ["a", "b", "c", "d"]
>>> list(itertools.zip_longest(x, y))
[(1, 'a'), (2, 'b'), (3, 'c'), (None, 'd')]

那么,就可以将`better_grouper`函数中的`zip`用`zip_longest()`替代了。

import itertools as it
 
def grouper(inputs, n, fillvalue=None):
    iters = [iter(inputs)] * n
    return it.zip_longest(*iters, fillvalue=fillvalue)

再跑一下,就是这样的结果了。

>>> nums = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> print(list(grouper(nums, 4)))
[(1, 2, 3, 4), (5, 6, 7, 8), (9, 10, None, None)]

------

《跟老齐学Python:轻松入门》