你正要写的那个循环,标准库里很可能已经有了。islice、chain、pairwise 和 groupby 能覆盖大部分情况,而 groupby 有一条规则绝对不能跳过。
itertools 在标准库里,不用安装任何东西。里面的工具都接收迭代器、返回迭代器,所以能和上一篇的生成器组合使用,而且一样是惰性的。
islice:从无穷无尽的东西里取一段
生成器不能切片。
g = (n for n in range(10))
try:
print(g[:5])
except TypeError as err:
print('TypeError:', err)
输出:
TypeError: 'generator' object is not subscriptable
islice 也是切片,只不过靠逐个索取,而不是靠索引。
from itertools import islice
def naturals():
n = 1
while True:
yield n
n += 1
print(list(islice(naturals(), 5)))
print(list(islice(naturals(), 10, 15)))
输出:
[1, 2, 3, 4, 5]
[11, 12, 13, 14, 15]
参数的写法和 range 一样:只给终止值,或者给起始值和终止值。它不能倒着数,因为迭代器没法倒着走。
chain:一个循环遍历好几样东西
from itertools import chain
a, b, c = [1, 2], (3, 4), range(5, 7)
print(list(chain(a, b, c)))
输出:
[1, 2, 3, 4, 5, 6]
列表、元组和 range,当作一个序列来遍历,而且不用先拼出一个合并后的列表。
如果手上是列表的列表,chain.from_iterable 可以展开一层:
from itertools import chain
rows = [['ada', 'grace'], ['alan'], ['edsger', 'barbara']]
print(list(chain.from_iterable(rows)))
输出:
['ada', 'grace', 'alan', 'edsger', 'barbara']
pairwise:每个元素和它前面那个
计算相邻读数之间的差,大多数人会写成带索引的循环,再顺手送上一个差一错误。
from itertools import pairwise
temps = [12, 14, 13, 18, 18, 21]
for a, b in pairwise(temps):
print(f"{a} -> {b} {b - a:+d}")
输出:
12 -> 14 +2
14 -> 13 -1
13 -> 18 +5
18 -> 18 +0
18 -> 21 +3
六个读数得到五对,正是你想要的数量。pairwise 需要 Python 3.10 或更高版本。
groupby,以及绝对不能跳过的规则
groupby 分组的是相邻的元素。它不会把整个序列里键相同的元素都收集到一起——每当键变了,它就开一个新组。
from itertools import groupby
people = [('ada', 'eng'), ('alan', 'math'), ('grace', 'eng'), ('emmy', 'math')]
for role, group in groupby(people, key=lambda p: p[1]):
print(role, [name for name, _ in group])
输出:
eng ['ada']
math ['alan']
eng ['grace']
math ['emmy']
进去两种角色,出来四个组。没有出错:这是文档里写明的行为,也正是大家认定 groupby 有问题的原因。
先按同一个键排序,它就会按你的本意工作:
from itertools import groupby
people = [('ada', 'eng'), ('alan', 'math'), ('grace', 'eng'), ('emmy', 'math')]
def by_role(p):
return p[1]
for role, group in groupby(sorted(people, key=by_role), key=by_role):
print(role, [name for name, _ in group])
输出:
eng ['ada', 'grace']
math ['alan', 'emmy']
同一个函数既传给 sorted,也传给 groupby。一旦两者对不上,你又会得到四个组的结果。
还有第二个坑。每个组都是建立在同一个底层序列上的迭代器,只在你移到下一组之前有效:
from itertools import groupby
people = [('ada', 'eng'), ('grace', 'eng'), ('alan', 'math')]
groups = list(groupby(people, key=lambda p: p[1]))
for role, group in groups:
print(role, list(group))
输出:
eng []
math []
list() 在读取任何内容之前就一路走到了最后一组,前面的组全被甩在了后面。要么在循环里当场消费每个组,要么边走边构建一个真正的字典。
如果只想计数,collections 那篇讲过的 Counter 比这些都简短。
count、cycle、repeat
三个无穷的迭代器。只有配上会停下的东西才能用——islice、break,或者和有限序列一起 zip。
from itertools import count, cycle, repeat, islice
print(list(islice(count(10, 5), 4)))
print(list(islice(cycle('ab'), 5)))
print(list(zip('abc', repeat(0))))
输出:
[10, 15, 20, 25]
['a', 'b', 'a', 'b', 'a']
[('a', 0), ('b', 0), ('c', 0)]
cycle 会把见过的每个元素都存一份副本,所以这三个里只有它的内存会增长。
combinations 和 product
两层嵌套循环,从此不用自己嵌套了。
from itertools import combinations, product
print(list(combinations('abc', 2)))
print(list(product([0, 1], repeat=2)))
输出:
[('a', 'b'), ('a', 'c'), ('b', 'c')]
[(0, 0), (0, 1), (1, 0), (1, 1)]
combinations 给出每个无序对,每对只出现一次。product 是嵌套循环能产生的所有组合,repeat=2 表示对同一个序列循环两层。
它们都返回迭代器
上面每个例子都包了一层 list(),原因只有一个:不包的话,你拿到的是一个对象,而不是值。
from itertools import chain
c = chain([1, 2], [3])
print(c)
print(list(c))
print(list(c))
输出:
<itertools.chain object at 0x7f9038d17e80>
[1, 2, 3]
[]
地址在你的机器上会不一样。空的第二行不会变——只能遍历一遍,和本系列里的其他东西一样。
要点
islice(it, n)是从无穷序列里取值的办法。切片语法对迭代器无效。chain把几个序列当作一个来遍历;chain.from_iterable把列表的列表展开一层。pairwise给出每个元素和它的前一个元素,而且数量不会算错。groupby分组的是连续的一段,不是相同的值。 先按同一个键排序,并在移到下一组之前消费完当前组。- 这里的一切都是惰性的,只能遍历一遍。真正需要值的时候,用
list()包起来。