Blog

Python itertools:那些不用再手写的循环

你正要写的那个循环,标准库里很可能已经有了。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 比这些都简短。

countcyclerepeat

三个无穷的迭代器。只有配上会停下的东西才能用——islicebreak,或者和有限序列一起 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 会把见过的每个元素都存一份副本,所以这三个里只有它的内存会增长。

combinationsproduct

两层嵌套循环,从此不用自己嵌套了。

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() 包起来。

这篇文章对你有帮助吗?

点一颗爱心来评分!

平均评分 0 / 5. 投票总数: 0

还没有人投票。来做第一个评分的人吧。