【Python系列专栏】第四十一篇 Python中常用内建模块（itertools）

itertools

简介

Python的内建模块 itertools 提供了非常有用的用于操作迭代对象的函数。我们首先看看 itertools 提供的几个“无限”迭代器：

---

count

count() 返回的是一个无限的迭代器，默认初始值为0，步长为1（按Python的传参规则，只传入一个参数时，传入的参数被视作初始值）

>>> import itertools
>>> natuals = itertools.count(1)
>>> for n in natuals:
...     print(n)
...
1
2
3
...

上述代码会打印出自然数序列，但问题是它根本停不下来，只能按 Ctrl+C 退出。

---

cycle

cycle() 会把传入的一个序列无限重复下去：

>>> import itertools
>>> cs = itertools.cycle('ABC') # 注意字符串也是序列的一种，还可以是列表、元组等
>>> for c in cs:
...     print(c)
...
'A'
'B'
'C'
'A'
'B'
'C'
...

同样停不下来。

---

repeat

repeat() 负责把一个元素无限重复下去，不过 repeat() 提供了第二个参数，用于限定重复的次数：

>>> ns = itertools.repeat('A', 3)
>>> for n in ns:
...     print(n)
...
A
A
A

---

takewhile

前面介绍了几种产生“无限”迭代器的方法，有没有办法对它们进行控制呢？有的~我们可以通过 takewhile()等函数，根据条件判断来截取出有限的序列：

>>> natuals = itertools.count(1)
>>> ns = itertools.takewhile(lambda x: x <= 10, natuals)
>>> list(ns)
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

这里我们只截取序列中小于等于10的数，所以迭代器产生的数不符合该条件时就会停止迭代，也就不会无限地排列下去了。

除了 takewhile() 之外，itertools 还提供了一些非常有用的迭代器操作函数，在后面几个小节中会进行介绍。

---

chain

chain() 可以把一组迭代对象串联起来，形成一个更大的迭代器：

>>> for c in itertools.chain('ABC', 'XYZ'):
...     print(c)
...
A
B
C
X
Y
Z

---

groupby

groupby() 可以把迭代器中相邻的重复元素挑出来放在一起：

>>> for key, group in itertools.groupby('AAABBBCCAAA'):
...     print(key, list(group))
...
A ['A', 'A', 'A']
B ['B', 'B', 'B']
C ['C', 'C']
A ['A', 'A', 'A']

实际上挑选规则是通过函数完成的，只要作用于函数的两个元素返回的值相等，这两个元素就被认为是同一组的，而函数返回值将作为该组的key。如果我们想忽略大小写来分组，可以让元素 'A' 和 'a' 都返回相同的key：

>>> for key, group in itertools.groupby('AaaBBbcCAAa', lambda c: c.upper()):
...     print(key, list(group))
...
A ['A', 'a', 'a']
B ['B', 'B', 'b']
C ['c', 'C']
A ['A', 'A', 'a']

---

小结

itertools 模块提供的全部是处理迭代功能的函数，它们的返回值不是 list，而是 Iterator，也即它们的值不是立刻计算出放在内存中的，只有进行迭代时（例如使用 for 循环）才会被真正计算出来。