如何使用Trie在已分词句子列表中搜索短语并返回匹配的完整句子

核心思路 #

由于输入是已经完成分词的语料，我们可以直接构建token级Trie（而非字符级Trie），每个Trie节点对应一个分词结果，节点中存储所有包含当前节点对应前缀短语的完整句子，查询时沿Trie匹配目标短语的所有token，匹配完成后直接取节点中存储的句子即可。

具体实现步骤 #

1. 定义Trie节点结构 #

每个节点包含两个核心字段：

• 子节点映射：key为分词token，value为对应子Trie节点
• 句子集合：存储所有包含从根节点到当前节点路径对应短语的完整句子，可选存储句子ID减少内存占用

2. 构建Trie索引 #

遍历所有已分词的句子，对每个句子枚举所有可能的短语起始位置，将从该位置开始的所有连续token序列插入Trie，每插入一个token就将当前完整句子存入对应节点的句子集合中。

3. 执行短语查询 #

将目标短语先做分词（若输入已分词可跳过该步骤），从Trie根节点开始依次匹配每个token：

• 任意token匹配失败直接返回空结果，说明无符合要求的句子
• 所有token匹配完成后，返回当前节点存储的所有句子即可

可运行代码示例（Python） #

from collections import defaultdict

class TrieNode:
    def __init__(self):
        # 子节点映射，key为分词token，value为子节点
        self.children = defaultdict(TrieNode)
        # 存储包含当前路径对应短语的句子，用set自动去重
        self.matched_sentences = set()

class PhraseSearchTrie:
    def __init__(self):
        self.root = TrieNode()
        # 可选：句子ID映射表，大语料下用ID代替原始句子存到节点中减少内存
        self.id_to_sent = dict()
        self.next_id = 0

    def add_sentence(self, tokenized_sent: list[str], raw_sent: str) -> None:
        """插入一条已分词的句子到Trie中"""
        # 大语料下可开启ID映射逻辑
        # sent_id = self.next_id
        # self.id_to_sent[sent_id] = raw_sent
        # self.next_id += 1

        sent_length = len(tokenized_sent)
        # 枚举所有短语起始位置
        for start_idx in range(sent_length):
            current_node = self.root
            # 插入从start_idx开始的所有连续token序列
            for end_idx in range(start_idx, sent_length):
                current_token = tokenized_sent[end_idx]
                current_node = current_node.children[current_token]
                # 存入当前句子，大语料下替换为sent_id
                current_node.matched_sentences.add(raw_sent)
    
    def search(self, tokenized_phrase: list[str]) -> list[str]:
        """查询包含目标短语的所有句子"""
        current_node = self.root
        for token in tokenized_phrase:
            if token not in current_node.children:
                return []
            current_node = current_node.children[token]
        # 大语料下这里替换为 [self.id_to_sent[id] for id in current_node.matched_sentences]
        return list(current_node.matched_sentences)

# 测试用例
if __name__ == "__main__":
    search_trie = PhraseSearchTrie()
    # 插入测试句子
    search_trie.add_sentence(["我", "喜欢", "吃", "苹果"], "我喜欢吃苹果")
    search_trie.add_sentence(["苹果", "富含", "维生素"], "苹果富含维生素")
    search_trie.add_sentence(["我", "明天", "买", "苹果"], "我明天买苹果")
    # 查询包含短语["吃", "苹果"]的句子
    print(search_trie.search(["吃", "苹果"])) # 输出: ['我喜欢吃苹果']
    # 查询包含短语["苹果"]的句子
    print(search_trie.search(["苹果"])) # 输出: ['我喜欢吃苹果', '苹果富含维生素', '我明天买苹果']

优化方案 #

• 语料规模超过10万条时，建议用句子ID替代原始句子存入Trie节点，内存占用可降低70%以上
• 若业务场景不需要支持超长短语查询，可限制插入Trie的短语最大长度（比如最多插入长度为15的短语），进一步降低内存消耗和构建耗时
• 支持模糊匹配的场景可在节点中增加容错逻辑，常规精确匹配场景无需修改

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