一种基于defaultdict的前缀树Python实现

杜逸先

发布于 2023-12-06 16:12:33

2680

发布于 2023-12-06 16:12:33

文章被收录于专栏：追不上乌龟的兔子

前缀树（Trie 树，也称为字典树、单词查找树）是一种树形数据结构，用于高效地存储和检索字符串集合中的键。前缀树的主要优势在于能够快速地查找具有相同前缀的字符串，并且对于大量的字符串集合，它可以提供较高的检索效率。

前缀树的应用非常广泛，包括：

字符串检索：通过前缀树可以快速查找是否存在某个字符串，或者查找具有相同前缀的所有字- 符串。
自动完成：前缀树可以用于实现自动完成功能，根据用户输入的前缀提供可能的建议。
IP 路由：在路由表中，前缀树用于快速匹配最长前缀。

前缀树可以通过多种方式实现，在 Python 中最简单且直观的方式是用嵌套的 dict 实现。

首先定义一下前缀树的接口，应该包括 insert、search 和 startswith 三个方法。

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")

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 属性的类型注解就很难写，一定要写话就会非常冗长。

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 属性来区分节点是否为单词结尾。

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 进行一些改造。

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 在很多场景下都是挺好用的。

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2023-12-05，如有侵权请联系 cloudcommunity@tencent.com 删除

python