collections 模块里有三个工具,可以替代你手写过无数次的循环。本文逐一介绍,并讲清楚 defaultdict 那个常让人意外的行为。
计数、分组和维护队列,是大多数人都手写过的三种循环。collections 模块三样都有现成的。
这里用到的全都在标准库里。
手动计数
用 dict.get 的写法最常见。
# counting the long way
words = ['apple', 'bread', 'apple', 'milk', 'bread', 'apple']
counts = {}
for w in words:
counts[w] = counts.get(w, 0) + 1
print(counts)
输出:
{'apple': 3, 'bread': 2, 'milk': 1}
Counter
Counter 接受任意可迭代对象并计数。访问不存在的键会返回 0,而不是抛出异常。
# Counter does it in one line
from collections import Counter, defaultdict, deque
c = Counter(words)
print(c)
print(c.most_common(2))
print(c['apple'], c['missing'])
输出:
Counter({'apple': 3, 'bread': 2, 'milk': 1})
[('apple', 3), ('bread', 2)]
3 0
大家用它,主要就是为了 most_common。自己按值给字典排序,得写一行代码再加一个 lambda。
Counter 的加减运算
Counter 可以相加和相减。相减时,结果为零或负数的项会被去掉。
# Counter arithmetic
stock = Counter(apple=5, bread=2)
sold = Counter(apple=3, bread=1, milk=1)
print(stock - sold)
print(stock + sold)
输出:
Counter({'apple': 2, 'bread': 1})
Counter({'apple': 8, 'bread': 3, 'milk': 1})
算库存、比较两次计数的差异、做简单的词袋统计,都很方便。
用 defaultdict 分组
defaultdict(list) 遇到不存在的键时会调用 list(),所以可以直接追加,不用先检查。
# grouping with defaultdict
people = [('eng', 'ada'), ('ops', 'grace'), ('eng', 'alan')]
groups = defaultdict(list)
for team, name in people:
groups[team].append(name)
print(dict(groups))
输出:
{'eng': ['ada', 'alan'], 'ops': ['grace']}
defaultdict 的坑
读取不存在的键,会把这个键创建出来。这就是让人意外的地方,通常表现为:跑完一个查找循环后,字典里的键比预想的多。
# the gotcha, reading creates the key
d = defaultdict(list)
print(d['not-there'])
print(dict(d))
plain = {}
print(plain.get('not-there'))
print(plain)
输出:
[]
{'not-there': []}
None
{}
普通字典配合 .get() 不会这样。如果只是想安全地读取,就用 .get()。
setdefault 也能做到
如果你不想导入任何东西,用 setdefault 也能达到目的。
# setdefault does the same without a factory
grouped = {}
for team, name in people:
grouped.setdefault(team, []).append(name)
print(grouped)
输出:
{'eng': ['ada', 'alan'], 'ops': ['grace']}
在密集循环里它会慢一些,因为每一轮都会新建空列表,不管用不用得上。
用 deque 做队列
从列表头部删除元素很慢,因为后面的所有元素都要往前挪。deque 在两端添加和删除都是常数时间。
# deque for queues, because list.pop(0) is slow
q = deque(['a', 'b', 'c'])
q.appendleft('start')
print(q.popleft(), list(q))
recent = deque(maxlen=3)
for item in range(6):
recent.append(item)
print(list(recent))
输出:
start ['a', 'b', 'c']
[3, 4, 5]
maxlen 提供了一个滑动窗口,会自动丢弃最旧的元素,很适合保存最近的记录和计算移动平均。
要点
Counter能对任意可迭代对象计数,还提供most_common。- Counter 支持
+和-。 defaultdict在读取时就会创建键,不只是写入时。- 需要从头部弹出元素时,
deque才是合适的类型。
这个模块里还有 namedtuple 和 ChainMap。把 collections 从头到尾读一遍,是你在标准库上能花的最值得的时间之一。