python的前缀树（字典树）

写在前面

说实话做题之前是真没听说过这个数据结构。

是啥

Trie，又称前缀树或字典树，是一棵有根树，其每个节点包含以下字段：

• 指向子节点的数组 children 。
• 布尔字段 isEnd，表示该节点是否为字符串的结尾。

插入字符串

我们从字典树的根开始，插入字符串。对于当前字符对应的子节点，有两种情况：

• 子节点存在。沿着指针移动到子节点，继续处理下一个字符。
• 子节点不存在。创建一个新的子节点，记录在 children 数组的对应位置上，然后沿着指针移动到子节点，继续搜索下一个字符。

重复以上步骤，直到处理字符串的最后一个字符，然后将当前节点标记为字符串的结尾。

查找前缀

我们从字典树的根开始，查找前缀。对于当前字符对应的子节点，有两种情况：

• 子节点存在。沿着指针移动到子节点，继续搜索下一个字符。
• 子节点不存在。说明字典树中不包含该前缀，返回空指针。

重复以上步骤，直到返回空指针或搜索完前缀的最后一个字符。

若搜索到了前缀的末尾，就说明字典树中存在该前缀。此外，若前缀末尾对应节点的 isEnd 为真，则说明字典树中存在该字符串。

实现

class Trie:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.next = [None] * 26
        self.isend = False


    def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        cur = self
        for c in word:
            od = ord(c) - ord('a')
            if cur.next[od] is None:
                cur.next[od] = Trie()
            cur = cur.next[od]
        cur.isend = True


    def search(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        cur = self
        for c in word:
            od = ord(c) - ord('a')
            if cur.next[od] is None:
                return False
            cur = cur.next[od]
        return cur.isend == True


    def startsWith(self, prefix: str) -> bool:
        """
        Returns if there is any word in the trie that starts with the given prefix.
        """
        cur = self
        for c in prefix:
            od = ord(c) - ord('a')
            if cur.next[od] is None:
                return False
            cur = cur.next[od]
        return True

如果节点不止26个小写字母需要把next数组替换成字典。

trie = {}
for prefix in dictionary:
    cur = trie
    for c in prefix:
        if c not in cur:
            cur[c] = {}
        cur = cur[c]
    cur['#'] = ''

不用类用纯字典也可以实现，这里’#'号表示结尾。

Trie = lambda :defaultdict(Trie)
trie = Trie()
for prefix in dictionary:
    cur = reduce(dict.__getitem__, prefix, trie)
    cur['#'] = ''

用defaultdict的话就更简短了。

实战

一般是先构建一个字典树，在干嘛干嘛，注意需要的话在构建的时候可以保留尾结点用于一些判断。