如何实现Trie树中匹配指定前缀的字符串查询功能？

嘿，我来帮你搞定这个Trie树的前缀搜索问题！先梳理下你现有代码的问题，然后一步步补全实现你要的功能。

首先，你的add_word方法还没写完，而且Node的value赋值逻辑存在漏洞。我们先把基础的单词添加功能补好，再实现前缀搜索的核心逻辑。

完整实现代码 #

这里提供两种实现思路，一种直观易懂，另一种更节省内存，你可以按需选择：

方式一：节点存储完整字符串（直观易读）

class Node:
    def __init__(self):
        self.children = [None] * 26
        self.end = False  # 标记当前节点是否为一个单词的结尾
        self.value = ""   # 存储到当前节点的完整字符串

class Trie:
    def __init__(self):
        self.root = Node()

    def add_word(self, key):
        current = self.root
        current_str = ""
        for char in key:
            # 计算字符在children数组中的索引
            position = ord(char) - ord('a')
            current_str += char
            # 字符对应节点不存在则创建
            if not current.children[position]:
                current.children[position] = Node()
            current = current.children[position]
            # 给当前节点赋值完整字符串
            current.value = current_str
        # 标记该节点为单词结尾
        current.end = True

    def search_prefix(self, prefix):
        current = self.root
        # 先定位到前缀的最后一个节点
        for char in prefix:
            position = ord(char) - ord('a')
            # 中途找不到节点，说明无匹配前缀，返回空列表
            if not current.children[position]:
                return []
            current = current.children[position]
        
        # 从前缀节点开始，深度优先搜索收集所有单词
        result = []
        self._dfs(current, result)
        return result

    # 辅助DFS函数，遍历子节点收集所有完整单词
    def _dfs(self, node, result):
        # 如果当前节点是单词结尾，加入结果列表
        if node.end:
            result.append(node.value)
        # 遍历所有子节点，递归搜索
        for child in node.children:
            if child:
                self._dfs(child, result)

方式二：动态拼接字符串（更省内存）

如果单词数量多、长度长，不想在节点中存储完整字符串，可以在遍历的时候动态拼接：

class Node:
    def __init__(self):
        self.children = [None] * 26
        self.end = False  # 仅标记是否为单词结尾

class Trie:
    def __init__(self):
        self.root = Node()

    def add_word(self, key):
        current = self.root
        for char in key:
            position = ord(char) - ord('a')
            if not current.children[position]:
                current.children[position] = Node()
            current = current.children[position]
        current.end = True

    def search_prefix(self, prefix):
        current = self.root
        for char in prefix:
            position = ord(char) - ord('a')
            if not current.children[position]:
                return []
            current = current.children[position]
        
        result = []
        self._dfs(current, prefix, result)
        return result

    def _dfs(self, node, current_str, result):
        if node.end:
            result.append(current_str)
        # 遍历26个字母对应的子节点
        for i in range(26):
            if node.children[i]:
                next_char = chr(ord('a') + i)
                self._dfs(node.children[i], current_str + next_char, result)

测试验证 #

用这段代码测试，完全符合你的预期：

trie = Trie()
trie.add_word("aa")
trie.add_word("aac")
trie.add_word("ab")

print(trie.search_prefix("aa"))  # 输出: ['aa', 'aac']
print(trie.search_prefix("a"))   # 输出: ['aa', 'aac', 'ab']
print(trie.search_prefix("b"))   # 输出: []

关键逻辑说明 #

1. 补全add_word：
   • 通过ord(char) - ord('a')计算字符在children数组中的索引
   • 不存在的节点必须创建，最后要给单词结尾节点的end标记为True
2. 前缀搜索核心：
   • 先遍历到前缀的最后一个节点，中途找不到直接返回空
   • 用DFS（深度优先搜索）遍历该节点的所有子树，收集所有标记为单词结尾的字符串——这就是你之前缺失的核心逻辑

内容的提问来源于stack exchange，提问作者clink