Python笔记：Trie树结构简介

• Python笔记：Trie树结构简介
  • 1. Trie树是什么
  • 2. Trie树原理
  • 3. Trie树代码实现
  • 4. Leetcode例题分析
    • 1. Leetcode 208. Implement Trie (Prefix Tree)
    • 2. Leetcode 211. Design Add and Search Words Data Structure
    • 3. Leetcode 1032. Stream of Characters
    • 4. Leetcode 212. Word Search II
  • 5. 参考链接

1. Trie树是什么

Trie树又名字典树，前序查找树等等。本质来说，他就是一个确定有限状态自动机（DFA），从根节点往下到叶子节点的每一条路径都是一次有效搜索。

通过Trie树，我们可以将搜索的时间复杂度大幅减小置近乎常数的量级。

2. Trie树原理

trie树的原理可以参考最后一章节中的参考链接里的几个博客说明，他们讲的都挺好的，配合下述原理图（图片来源于网上），相信理解Trie树的原理将完全不会有任何问题。

下面，我们直接来考察Trie树的代码实现。

3. Trie树代码实现

Trie树的代码实现事实上也不复杂，尤其是在使用python的情况下。他的本质事实上就是一个用字典构建的树。

trie树最核心的功能包括以下两点：

1. 插入新的词汇；
2. 从字典中查询词汇；

下面，给出最基本的代码实现如下：

class Trie:
    def __init__(self):
        self.trie = {
     }
    
    def add_word(self, word):
        trie = self.trie
        for c in word:
            trie = trie.setdefault(c, {
     })
        trie["eos"] = ""

    def find(self, word):
        trie = self.trie
        for c in word:
            if c not in trie:
                return False
            trie = trie[c]
        return "eos" in trie

当然，上述只是最为简单的trie树功能实现，实际在使用中会存在多种变体，比如：

• 纯前缀查找（不完全匹配）
• 某前缀下的所有可能搜索结果
• 给字典中每一个单词加入权重
• ……

下面，我们来结合leetcode上面的习题来对trie树进行一些更深入的讨论。

4. Leetcode例题分析

leetcode中有一个专门的Trie树的题目集合，目前有17道题目，其中3题被锁住了，剩余14题是全员可见的，拿来练手足够了。

我们从中选取几道题目来浅谈一下Trie树的一些常规用法。

相信如果把下面几道Leetcode中的Trie树的题目搞定的话，你对Trie树应该至少也有一个非常直观的理解了。

1. Leetcode 208. Implement Trie (Prefix Tree)

给出题目链接如下：https://leetcode.com/problems/implement-trie-prefix-tree/

这一题算是Trie树的最基础例题了，他事实上就是要求你实现一下Trie树，并完成其中的完全匹配搜索与不完全匹配搜索。

我们直接给出代码实现如下：

class Trie:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.cache = {
     }
        

    def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        cache = self.cache
        for c in word:
            cache = cache.setdefault(c, {
     })
        cache["eos"] = {
     }
        return

    def search(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        cache = self.cache
        for c in word:
            if c not in cache.keys():
                return False
            cache = cache[c]
        return "eos" in cache.keys()
        

    def startsWith(self, prefix: str) -> bool:
        """
        Returns if there is any word in the trie that starts with the given prefix.
        """
        cache = self.cache
        for c in prefix:
            if c not in cache.keys():
                return False
            cache = cache[c]
        return True

有了这部分的基础，我们后续就可以进行更多的实例考察了。

2. Leetcode 211. Design Add and Search Words Data Structure

给出题目链接如下：https://leetcode.com/problems/design-add-and-search-words-data-structure/

这一题是上一题的变体，他的难度会更高一点，要求是允许.匹配任意一个字符，因此，每当遇到一次.，我们就需要遍历当前根下的所有节点，看是否存在正确的匹配。

本质上来说，就是要我们在匹配完成前序数组之后遍历Trie树中在该前缀条件下的所有可能的完全匹配结果。其中，前者我们通过trie树可以轻松实现，后者我们可以通过深度优先算法（dfs）来进行实现。

给出代码实现如下：

class WordDictionary:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.trie = {
     }

    def addWord(self, word: str) -> None:
        """
        Adds a word into the data structure.
        """
        trie = self.trie
        for c in word:
            trie = trie.setdefault(c, {
     })
        trie["eos"] = ""
        return

    def search(self, word: str) -> bool:
        """
        Returns if the word is in the data structure. A word could contain the dot character '.' to represent any one letter.
        """
        n = len(word)
        def dfs(trie, i):
            if i == n:
                return "eos" in trie
            if word[i] != '.':
                if word[i] not in trie:
                    return False
                else:
                    return dfs(trie[word[i]], i+1)
            else:
                return any(dfs(trie[c], i+1) for c in trie.keys() if c != "eos")
            
        return dfs(self.trie, 0)

这一题本质上就是需要在匹配了部分前序字符串之后遍历所有的可能匹配，这部分内容我们都可以通过Trie树加上dfs算法来实现。

类似的leetcode题目还包括：

1. 677. Map Sum Pairs
2. 676. Implement Magic Dictionary
3. 336. Palindrome Pairs

其中，前两题和本题基本是完全一致的，仿照着做就行了，第三题多少稍微复杂点，但是一旦确定思路就是反向找到所有不完全匹配的字符串集合之后再考察其是否组合结果是否满足回文条件的话，那么本质上来说就又回到这一题的算法了，没有根本性的差别。

有兴趣的读者可以自行尝试去做一下，这里就不再赘述了。

3. Leetcode 1032. Stream of Characters

给出题目链接如下：https://leetcode.com/problems/stream-of-characters/

这一题事实上挺简单的，就是将前序查找变换后后续查找，我们在Trie树的元素加入过程中直接反序地将元素加入之后即可。

给出代码实现如下：

class Trie:
    def __init__(self):
        self.trie = {
     }

    def add(self, word):
        trie = self.trie
        for c in word:
            trie = trie.setdefault(c, {
     })
        trie["eos"] = ""
    
    def find(self, word):
        trie = self.trie
        for c in word:
            if "eos" in trie:
                return True
            if c not in trie:
                return False
            trie = trie[c]
        return "eos" in trie

class StreamChecker:

    def __init__(self, words: List[str]):
        self.querys = ""
        self.trie = Trie()
        for word in words:
            self.trie.add(word[::-1])

    def query(self, letter: str) -> bool:
        self.querys = letter + self.querys
        return self.trie.find(self.querys)

4. Leetcode 212. Word Search II

给出题目链接如下：https://leetcode.com/problems/word-search-ii/

这一题就是基于Trie树的一道纯应用题了，用的就是最基本的Trie树结果，但是需要结合栈的内容来实现邻接序列的遍历。

给出代码实现如下：

class Trie:
    def __init__(self, words):
        self.trie = {
     }
        for word in words:
            self.add(word)
            
    def add(self, word):
        trie = self.trie
        for c in word:
            trie = trie.setdefault(c, {
     })
        trie["eos"] = "eos"
        
    def find(self, word, mode="full"):
        trie = self.trie
        for c in word:
            if c not in trie:
                return False
            trie = trie[c]
        if mode != "full":
            return True
        else:
            return "eos" in trie

class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        have_visited = set()
        ans = set()
        stack = []
        n = len(board)
        m = len(board[0])
        trie = Trie(words)
        
        def dfs(row, col):
            nonlocal stack, have_visited, ans
            stack.append(board[row][col])
            have_visited.add((row, col))
            if not trie.find(stack, mode='half'):
                stack.pop()
                have_visited.remove((row, col))
                return
            if trie.find(stack, mode='full'):
                ans.add(''.join(stack))
            if row-1 >= 0 and (row-1, col) not in have_visited:
                dfs(row-1, col)
            if row+1 < n and (row+1, col) not in have_visited:
                dfs(row+1, col)
            if col-1 >= 0 and (row, col-1) not in have_visited:
                dfs(row, col-1)
            if col+1 < m and (row, col+1) not in have_visited:
                dfs(row, col+1)
            stack.pop()
            have_visited.remove((row, col))
            return
        
        for i in range(n):
            for j in range(m):
                dfs(i, j)
        return list(ans)

5. 参考链接

1. WikiPedia: Trie树
2. 知乎：浅谈trie树
3. Trie树（Prefix Tree）介绍