前缀树（Trie 树，也称为字典树、单词查找树）是一种树形数据结构，用于高效地存储和检索字符串集合中的键。前缀树的主要优势在于能够快速地查找具有相同前缀的字符串，并且对于大量的字符串集合，它可以提供较高的检索效率。
前缀树的应用非常广泛，包括：
- 字符串检索：通过前缀树可以快速查找是否存在某个字符串，或者查找具有相同前缀的所有字- 符串。
- 自动完成：前缀树可以用于实现自动完成功能，根据用户输入的前缀提供可能的建议。
- IP 路由：在路由表中，前缀树用于快速匹配最长前缀。
前缀树可以通过多种方式实现，在 Python 中最简单且直观的方式是用嵌套的 dict 实现。
首先定义一下前缀树的接口，应该包括 insert、search 和 startswith 三个方法。
1
2
3
4
5
6
7
8
9
10
| trie = Trie()
trie.insert("app")
trie.insert("apple")
trie.insert("banana")
assert trie.startswith("a")
assert trie.startswith("ba")
assert not trie.startswith("c")
assert trie.search("app")
assert not trie.search("appl")
assert trie.search("apple")
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
| class Trie:
def __init__(self):
self.root = {-1: False}
def insert(self, word: str):
parent = self.root
for ch in word:
node = parent.get(ch, {-1: False})
parent[ch] = node
parent = node
parent[-1] = True
def __find(self, word:str):
node = self.root
for ch in word:
if ch not in node:
return None
node = node[ch]
return node
def search(self, word: str) -> bool:
node = self.__find(word)
return node is not None and node[-1]
def startswith(self, prefix: str) -> bool:
node = self.__find(prefix)
return node is not None
|
这个实现总体上比较简单，每一个节点都是一个 dict，key 是字符，value 是下一个节点。为了区分一个节点是否是一个单词的结尾，我们使用 -1 作为特殊的 key，用于存储该节点是否是一个单词的结尾。
不过这个实现有一些小问题。首先用 dict 承担了不同的职责，既用于存储下一个节点，又用于存储是否是单词结尾的标记，语义上不是很清晰，或者说 root 属性的类型注解就很难写，一定要写话就会非常冗长。
1
2
3
4
5
6
7
| from typing import TypeAlias, Literal, Union
TrieNode: TypeAlias = dict[str | Literal[-1], Union[bool, "TrieNode"]]
class Trie:
def __init__(self):
self.root: TrieNode = {-1: False}
|
另外这个实现多次进行了节点的初始化，分别在初始化对象时 self.root: TrieNode = {-1: False}和 insert 方法中 node = parent.get(ch, {-1: False})。python 标准库中的defaultdict可以帮助我们简化这种重复初始化 dict 的值的操作。
我们可以重新设计一下 TrieNode 的类型，使用defaultdict来初始化节点，并通过单独的 end 属性来区分节点是否为单词结尾。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
| from collections import defaultdict
class TrieNode(defaultdict):
def __init__(self):
super().__init__(TrieNode)
self.end = False
class Trie:
def __init__(self) -> None:
self.root = TrieNode()
def insert(self, word: str) -> None:
node = self.root
for ch in word:
node = node[ch]
node.end = True
def __find(self, word: str) -> TrieNode:
node = self.root
for ch in word:
if ch not in node:
return None
node = node[ch]
return node
def startswith(self, prefix: str) -> bool:
return self.__find(prefix) is not None
def search(self, word: str) -> bool:
return (node := self.__find(word)) is not None and node.end
|
TrieNode 类继承自 defaultdict，这样我们就可以通过node = node[ch]来初始化节点，独立的 end 属性也可以简化节点是否为单词结尾的判断，比起使用特殊 key 的方式语义上更加清晰。
上述的两个 Trie 实现都可以用于字符串（前缀）检索场景，如果我们想要实现自动完成功能，就需要对 Trie 进行一些改造。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
| class TrieNode(defaultdict):
...
def iter_words(self, prefix=""):
if self.end:
yield prefix
for ch, child in self.items():
yield from child.iter_words(prefix + ch)
class Trie:
...
def words_with_prefix(self, prefix: str) -> str:
node = self.__find(prefix)
return list(node.iter_words(prefix)) if node else []
|
首先为 TrieNode 添加一个iter_words方法，用于遍历以当前节点为根节点的所有单词。
然后在Trie类中添加words_with_prefix方法，用于返回以prefix为前缀的所有单词。
P.S. 这个继承defaultdict的方式实现 TrieNode 的方法实际上来自 Github Copilot 的代码改进建议。Copilot 在很多场景下都是挺好用的。
