Python 实例｜前缀树（Trie）

我们实现一个简单的前缀树，支持以下 3 个方法：

• 支持使用 $O(S)$ 的时间复杂度，将一个词语添加到前缀树中，其中 $S$ 为字符串长度
• 支持使用 $O(S)$ 的时间复杂度，判断一个字符串是否在前缀树中，其中 $S$ 为字符串长度
• 支持使用 $O(S\times T)$ 的时间复杂度，获取一个字符串中所有存在于前缀树中的子串，其中 $S$ 为字符串长度，$T$ 为前缀树深度

from typing import List


class Trie:
    """前缀树

    支持以下 3 个方法：
    - 支持使用 O(S) 的时间复杂度，将一个词语添加到前缀树中，其中 S 为字符串长度
    - 支持使用 O(S) 的时间复杂度，判断一个字符串是否在前缀树中，其中 S 为字符串长度
    - 支持使用 O(S·T) 的时间复杂度，获取一个字符串中所有存在于前缀树中的子串，其中 S 为字符串长度，T 为前缀树深度
    """

    def __init__(self):
        self._trie = {
    
    }  # 初始化前缀树的根节点

    def add_word(self, word: str) -> None:
        """将一个词语 word 添加到前缀树中

        Parameters
        ----------
        word : str
            需要添加的词语
        """
        node = self._trie
        for ch in word:
            if ch not in node:
                node[ch] = {
    
    }
            node = node[ch]
        node["$"] = word  # 标记当前位置可以是完整词语

    def match(self, s: str) -> bool:
        """判断一个字符串 s 是否在前缀树中

        Parameters
        ----------
        s : str
            需要查询的字符串

        Returns
        -------
        bool
            字符串 s 是否在前缀树中
        """
        node = self._trie
        for ch in s:
            if ch not in node:
                return False
            node = node[ch]
        return "$" in node

    def findall(self, s: str) -> List[str]:
        """获取一个字符串 s 中所有存在于前缀树中的子串

        Parameters
        ----------
        s : str
            需要查询的字符串

        Returns
        -------
        List[str]
            返回字符串 s 中所有存在于前缀树中的子串
        """
        res = []
        for i in range(len(s)):
            node = self._trie
            for j in range(i, len(s)):
                ch = s[j]
                if ch not in node:
                    break
                node = node[ch]
                if "$" in node:
                    res.append(node["$"])
        return res