Trie tree 解题模板

Posted on Edited on In ,

什么是 Trie 树？

Trie 树（发音”try”），又称前缀树或字典树，是一种树形数据结构，用于高效存储和检索字符串数据集中的键。

典型应用场景:

  • 自动补全（搜索框）
  • 拼写检查
  • IP 路由（最长前缀匹配）
  • 单词搜索游戏
  • 字符串前缀/后缀查询

时间复杂度:

  • 插入: O(m) - m 是字符串长度
  • 查找: O(m)
  • 前缀搜索: O(m)

优势:

  • 查找效率高于哈希表（处理前缀）
  • 不存在哈希冲突
  • 可以按字典序遍历

模板 (Python)

模板 1: 标准 Trie 树实现

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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
class TrieNode:
    """Trie树的节点"""
    def __init__(self):
        # 子节点字典 {'a': TrieNode, 'b': TrieNode, ...}
        self.children = {}
        # 标记是否为单词结尾
        self.is_end_of_word = False
        # 可选：存储单词本身（某些题目需要）
        # self.word = None


class Trie:
    """标准Trie树实现"""

    def __init__(self):
        """初始化根节点"""
        self.root = TrieNode()


    def insert(self, word: str) -> None:
        """
        插入单词到Trie树
        时间复杂度: O(m), m为单词长度
        """
        node = self.root

        for char in word:
            # 如果字符不存在，创建新节点
            if char not in node.children:
                node.children[char] = TrieNode()
            # 移动到下一个节点
            node = node.children[char]

        # 标记单词结束
        node.is_end_of_word = True
        # node.word = word  # 可选：存储完整单词


    def search(self, word: str) -> bool:
        """
        查找完整单词是否存在
        时间复杂度: O(m)
        """
        node = self.root

        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]

        # 必须是完整单词
        return node.is_end_of_word


    def starts_with(self, prefix: str) -> bool:
        """
        查找是否存在以prefix开头的单词
        时间复杂度: O(m)
        """
        node = self.root

        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]

        # 只要前缀存在即可，不需要是完整单词
        return True


    def delete(self, word: str) -> bool:
        """
        删除单词（可选功能）
        返回是否成功删除
        """
        def _delete_helper(node, word, index):
            if index == len(word):
                # 到达单词末尾
                if not node.is_end_of_word:
                    return False  # 单词不存在

                node.is_end_of_word = False
                # 如果该节点没有子节点，可以删除
                return len(node.children) == 0

            char = word[index]
            if char not in node.children:
                return False  # 单词不存在

            child = node.children[char]
            should_delete_child = _delete_helper(child, word, index + 1)

            if should_delete_child:
                del node.children[char]
                # 如果当前节点不是其他单词的结尾且没有其他子节点
                return not node.is_end_of_word and len(node.children) == 0

            return False

        return _delete_helper(self.root, word, 0)


# 使用示例
trie = Trie()
trie.insert("apple")
print(trie.search("apple"))      # True
print(trie.search("app"))        # False
print(trie.starts_with("app"))   # True
trie.insert("app")
print(trie.search("app"))        # True

模板 2: 紧凑版 Trie

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:
    """紧凑版Trie - 面试推荐"""

    def __init__(self):
        self.root = {}

    def insert(self, word: str) -> None:
        node = self.root
        for char in word:
            node = node.setdefault(char, {})
        node['#'] = True  # 用'#'标记单词结束

    def search(self, word: str) -> bool:
        node = self.root
        for char in word:
            if char not in node:
                return False
            node = node[char]
        return '#' in node

    def starts_with(self, prefix: str) -> bool:
        node = self.root
        for char in prefix:
            if char not in node:
                return False
            node = node[char]
        return True

模板 3: 带通配符的 Trie (支持’.’)

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
33
34
35
36
37
38
39
40
41
42
43
44
45
class WordDictionary:
    """支持'.'通配符的单词字典"""

    def __init__(self):
        self.root = {}

    def addWord(self, word: str) -> None:
        node = self.root
        for char in word:
            node = node.setdefault(char, {})
        node['#'] = True

    def search(self, word: str) -> bool:
        """支持'.'匹配任意单个字符"""
        return self._search_helper(word, 0, self.root)

    def _search_helper(self, word: str, index: int, node: dict) -> bool:
        # 递归终止条件
        if index == len(word):
            return '#' in node

        char = word[index]

        if char == '.':
            # '.'可以匹配任意字符
            for child in node:
                if child != '#' and self._search_helper(word, index + 1, node[child]):
                    return True
            return False
        else:
            # 普通字符
            if char not in node:
                return False
            return self._search_helper(word, index + 1, node[char])


# 使用示例
wd = WordDictionary()
wd.addWord("bad")
wd.addWord("dad")
wd.addWord("mad")
print(wd.search("pad"))    # False
print(wd.search("bad"))    # True
print(wd.search(".ad"))    # True
print(wd.search("b.."))    # True

LeetCode 经典题目与模板应用

题目 1: [208] Implement Trie (Prefix Tree)

难度: Medium

题意: 实现 Trie 的基本操作

解答: 使用标准模板即可

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
class Trie:
    def __init__(self):
        self.root = {}

    def insert(self, word: str) -> None:
        node = self.root
        for char in word:
            node = node.setdefault(char, {})
        node['#'] = True

    def search(self, word: str) -> bool:
        node = self.root
        for char in word:
            if char not in node:
                return False
            node = node[char]
        return '#' in node

    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for char in prefix:
            if char not in node:
                return False
            node = node[char]
        return True

题目 2: [211] Design Add and Search Words Data Structure

难度: Medium

题意: 支持’.’通配符的单词搜索

关键点:

  • ‘.’需要递归搜索所有可能的子节点
  • 使用 DFS 遍历

解答: 使用模板 3

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
class WordDictionary:
    def __init__(self):
        self.root = {}

    def addWord(self, word: str) -> None:
        node = self.root
        for char in word:
            node = node.setdefault(char, {})
        node['#'] = True

    def search(self, word: str) -> bool:
        def dfs(index, node):
            if index == len(word):
                return '#' in node

            char = word[index]
            if char == '.':
                # 尝试所有子节点
                for child in node:
                    if child != '#' and dfs(index + 1, node[child]):
                        return True
                return False
            else:
                if char not in node:
                    return False
                return dfs(index + 1, node[char])

        return dfs(0, self.root)

题目 3: [212] Word Search II

难度: Hard

题意: 在二维字符网格中搜索多个单词

关键点:

  • 将所有单词插入 Trie
  • 使用 DFS + Trie 剪枝
  • 找到单词后需要去重

完整解答:

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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
class Solution:
    def findWords(self, board: List[List[str]], words: List[str]) -> List[str]:
        # Step 1: 构建Trie树
        trie = {}
        for word in words:
            node = trie
            for char in word:
                node = node.setdefault(char, {})
            node['#'] = word  # 存储完整单词

        # Step 2: DFS搜索
        rows, cols = len(board), len(board[0])
        result = []
        visited = set()

        def dfs(r, c, node, path):
            # 边界检查
            if (r < 0 or r >= rows or c < 0 or c >= cols or
                (r, c) in visited or board[r][c] not in node):
                return

            char = board[r][c]
            node = node[char]

            # 找到单词
            if '#' in node:
                result.append(node['#'])
                del node['#']  # 避免重复添加

            # 标记访问
            visited.add((r, c))

            # 四个方向DFS
            for dr, dc in [(0,1), (1,0), (0,-1), (-1,0)]:
                dfs(r + dr, c + dc, node, path + char)

            # 回溯
            visited.remove((r, c))

        # Step 3: 从每个位置开始搜索
        for r in range(rows):
            for c in range(cols):
                dfs(r, c, trie, "")

        return result


# 时间复杂度: O(M * N * 4^L)
# M, N为网格大小，L为最长单词长度
# 空间复杂度: O(K * L)，K为单词数量

优化技巧:

1
2
3
4
5
6
7
# 优化1: 剪枝 - 如果某个节点下没有子节点了，删除它
if not node[char]:
    del node[char]

# 优化2: 提前终止 - 如果找到所有单词就停止
if not trie:
    return result

题目 4: [648] Replace Words

难度: Medium

题意: 用字典中的词根替换句子中的单词

示例:

1
2
Input: dictionary = ["cat","bat","rat"], sentence = "the cattle was rattled by the battery"
Output: "the cat was rat by the bat"

解答:

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
class Solution:
    def replaceWords(self, dictionary: List[str], sentence: str) -> str:
        # 构建Trie
        trie = {}
        for word in dictionary:
            node = trie
            for char in word:
                node = node.setdefault(char, {})
            node['#'] = word  # 存储词根

        # 查找最短前缀
        def find_root(word):
            node = trie
            for char in word:
                if char not in node:
                    return word  # 没有找到词根
                node = node[char]
                if '#' in node:
                    return node['#']  # 找到最短词根
            return word  # 完整单词没有词根

        # 处理句子
        words = sentence.split()
        return ' '.join(find_root(word) for word in words)


# 时间复杂度: O(N) - N为所有字符总数
# 空间复杂度: O(M) - M为字典中所有字符总数

题目 5: [677] Map Sum Pairs

难度: Medium

题意: 实现一个支持前缀和查询的数据结构

关键点: 在 Trie 节点中存储数值

解答:

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
33
34
35
class MapSum:
    def __init__(self):
        self.root = {}
        self.map = {}  # 存储key-value对，用于更新

    def insert(self, key: str, val: int) -> None:
        # 计算差值（如果key已存在）
        delta = val - self.map.get(key, 0)
        self.map[key] = val

        # 更新Trie中的值
        node = self.root
        for char in key:
            node = node.setdefault(char, {})
            node.setdefault('val', 0)
            node['val'] += delta  # 累加差值
        node['#'] = True

    def sum(self, prefix: str) -> int:
        node = self.root
        for char in prefix:
            if char not in node:
                return 0
            node = node[char]
        return node.get('val', 0)


# 使用示例
ms = MapSum()
ms.insert("apple", 3)
print(ms.sum("ap"))      # 3
ms.insert("app", 2)
print(ms.sum("ap"))      # 5
ms.insert("apple", 5)    # 更新apple的值
print(ms.sum("ap"))      # 7

题目 6: [720] Longest Word in Dictionary

难度: Medium

题意: 找出字典中最长的可以逐个字母构建的单词

示例:

1
2
3
Input: words = ["w","wo","wor","worl","world"]
Output: "world"
解释: "world"可以通过"w"->"wo"->"wor"->"worl"->"world"逐步构建

解答:

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
33
34
35
36
37
class Solution:
    def longestWord(self, words: List[str]) -> str:
        # 构建Trie
        trie = {}
        for word in words:
            node = trie
            for char in word:
                node = node.setdefault(char, {})
            node['#'] = word

        # DFS查找最长单词
        result = ""

        def dfs(node, path):
            nonlocal result

            # 遍历所有子节点
            for char in node:
                if char == '#':
                    continue

                child = node[char]
                # 必须是完整单词才能继续
                if '#' in child:
                    new_path = child['#']
                    # 更新结果（更长或字典序更小）
                    if len(new_path) > len(result) or \
                       (len(new_path) == len(result) and new_path < result):
                        result = new_path
                    dfs(child, new_path)

        dfs(trie, "")
        return result


# 时间复杂度: O(N) - N为所有字符总数
# 空间复杂度: O(N)

题目 7: [1268] Search Suggestions System

难度: Medium

题意: 实现搜索建议系统，每输入一个字符返回最多 3 个建议

示例:

1
2
3
4
5
6
7
8
9
Input: products = ["mobile","mouse","moneypot","monitor","mousepad"],
       searchWord = "mouse"
Output: [
    ["mobile","moneypot","monitor"],
    ["mobile","moneypot","monitor"],
    ["mouse","mousepad"],
    ["mouse","mousepad"],
    ["mouse","mousepad"]
]

解答:

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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        # 方法1: Trie + DFS
        # 构建Trie
        trie = {}
        for product in products:
            node = trie
            for char in product:
                node = node.setdefault(char, {})
            node['#'] = product

        result = []
        node = trie
        prefix = ""

        # 对每个前缀查找建议
        for char in searchWord:
            prefix += char

            if char not in node:
                # 没有匹配的产品
                result.append([])
                node = {}  # 标记为空，后续都返回空列表
            else:
                node = node[char]
                # DFS收集所有单词
                suggestions = []

                def dfs(n):
                    if len(suggestions) >= 3:  # 只需要3个
                        return
                    if '#' in n:
                        suggestions.append(n['#'])
                    # 按字典序遍历
                    for key in sorted(n.keys()):
                        if key != '#' and len(suggestions) < 3:
                            dfs(n[key])

                dfs(node)
                result.append(suggestions)

        return result


# 方法2: 排序 + 二分 (更简单)
class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        products.sort()  # 排序
        result = []
        prefix = ""

        for char in searchWord:
            prefix += char
            # 找到第一个匹配的位置
            i = bisect.bisect_left(products, prefix)
            # 取最多3个匹配的产品
            suggestions = []
            for j in range(i, min(i + 3, len(products))):
                if products[j].startswith(prefix):
                    suggestions.append(products[j])
            result.append(suggestions)

        return result

# Trie方法时间: O(N*M + K*M*log26) N个产品，M平均长度，K为searchWord长度
# 排序方法时间: O(N*logN + K*N) 更简单实用