验证程序经常需要做四件事:统计出现次数、按照某个键分组、判断状态是否访问过、从队首取出 BFS 状态。普通 dict、set 和 collections 中的 Counter、defaultdict、deque 可以直接表达这些操作。
普通字典 dict
字典相当于 C++ 的 unordered_map,保存键到值的映射:
1score = {"Alice": 95, "Bob": 90}
2
3assert score["Alice"] == 95
4assert "Carol" not in score
手动统计频率时,dict.get(key, default) 很方便:
1a = [3, 1, 3, 2, 3, 2]
2count = {}
3
4for x in a:
5 count[x] = count.get(x, 0) + 1
6
7assert count == {1: 1, 2: 2, 3: 3}
count.get(x, 0) 在键不存在时返回 0,但不会自动把这个键插入字典。
Counter:直接统计频率
1from collections import Counter
2
3a = [3, 1, 3, 2, 3, 2]
4count = Counter(a)
5
6assert count[3] == 3
7assert count[2] == 2
8assert count[100] == 0
9assert count.most_common(2) == [(3, 3), (2, 2)]
Counter 对不存在的键返回 0。most_common(k) 返回出现次数最多的前 k 项。
验证两个序列是否为同一多重集合
集合会丢失重复次数:
1assert set([1, 1, 2]) == set([1, 2, 2])
使用 Counter 才能同时比较元素和出现次数:
1from collections import Counter
2
3
4def is_permutation(source, answer):
5 return Counter(source) == Counter(answer)
6
7
8assert is_permutation([1, 1, 2], [2, 1, 1])
9assert not is_permutation([1, 1, 2], [1, 2, 2])
这适合验证“输出必须是输入的一个排列”之类的构造题条件。
Counter 运算
1from collections import Counter
2
3left = Counter("aab")
4right = Counter("abb")
5
6assert left + right == Counter({"a": 3, "b": 3})
7assert left - right == Counter({"a": 1})
8assert left & right == Counter({"a": 1, "b": 1})
9assert left | right == Counter({"a": 2, "b": 2})
+ 合并计数,- 只保留正计数,& 取每项最小值,| 取每项最大值。需要保留零或负计数时,不要依赖减法结果,改用 subtract 或普通字典检查。
defaultdict:自动创建默认值
按键分组时,普通字典需要先判断键是否存在;defaultdict 可以自动创建空容器:
1from collections import defaultdict
2
3groups = defaultdict(list)
4
5for x in [1, 2, 3, 4, 5, 6]:
6 groups[x % 3].append(x)
7
8assert groups[0] == [3, 6]
9assert groups[1] == [1, 4]
10assert groups[2] == [2, 5]
参数不是一个默认值,而是创建默认值的函数:
1from collections import defaultdict
2
3integer_count = defaultdict(int)
4sets_by_key = defaultdict(set)
5
6integer_count["x"] += 1
7sets_by_key[0].add(7)
8
9assert integer_count["x"] == 1
10assert sets_by_key[0] == {7}
构造邻接表
1from collections import defaultdict
2
3edges = [(1, 2), (1, 3), (2, 4)]
4graph = defaultdict(list)
5
6for u, v in edges:
7 graph[u].append(v)
8 graph[v].append(u)
9
10assert graph[1] == [2, 3]
11assert graph[4] == [2]
读取一个从未出现的键会把它插入 defaultdict。只想查询而不改变字典时,可以使用 key in graph 或 graph.get(key)。
set:判重与集合关系
1seen = set()
2
3for x in [3, 1, 3, 2]:
4 seen.add(x)
5
6assert seen == {1, 2, 3}
7assert 2 in seen
常用集合运算:
1a = {1, 2, 3}
2b = {3, 4}
3
4assert a | b == {1, 2, 3, 4}
5assert a & b == {3}
6assert a - b == {1, 2}
7assert a ^ b == {1, 2, 4}
集合元素必须可哈希。列表不能直接放进集合,状态通常转成元组:
1visited = set()
2state = [1, 2, 3]
3
4visited.add(tuple(state))
5
6assert (1, 2, 3) in visited
deque:双端队列
列表尾部的 append、pop 很快,但 pop(0) 需要移动后面所有元素。BFS 应使用 deque.popleft():
1from collections import deque
2
3queue = deque([2, 3])
4queue.append(4)
5queue.appendleft(1)
6
7assert queue.popleft() == 1
8assert queue.pop() == 4
9assert list(queue) == [2, 3]
BFS 模板
1from collections import deque
2
3
4def shortest_steps(start, target):
5 queue = deque([(start, 0)])
6 visited = {start}
7
8 while queue:
9 x, distance = queue.popleft()
10 if x == target:
11 return distance
12
13 for next_x in (x - 1, x + 1, x * 2):
14 if 0 <= next_x <= 100 and next_x not in visited:
15 visited.add(next_x)
16 queue.append((next_x, distance + 1))
17
18 return None
19
20
21assert shortest_steps(5, 17) == 4
22assert shortest_steps(7, 7) == 0
怎样选择
| 任务 | 数据结构 |
|---|---|
| 键到值的映射 | dict |
| 统计出现次数 | Counter |
| 按键收集列表或集合 | defaultdict |
| 只判断是否出现 | set |
| BFS 队列、双端操作 | deque |
| 可哈希的复合状态 | tuple |
容器应该表达题意,而不是为了使用标准库而使用。只有计数时用 Counter,需要附加信息时普通 dict 往往更清楚。
状态搜索的完整例子见用 Python 快速编写算法暴力验证程序,无序答案的排序比较见Python 排序与顺序验证。