如何从Trie结构中检索包含/不包含指定字母的单词

基于提供的Trie结构，可以通过深度优先遍历加路径剪枝的方式，高效实现「包含指定字母」和「不包含指定字母」的检索能力，不需要提前导出Trie中存储的全量词表。

已有基础代码 #

def load_trie(words):
    root = {}
    for word in words:
        curr_node = root
        for letter in word:
            curr_node = curr_node.setdefault(letter, {})
        curr_node.setdefault('', True)
    return root

with open('sowpods') as word_list:
    words = [word.strip().upper() for word in word_list]
    
TRIE = load_trie(words)

检索函数实现 #

def search_words(trie_root, required_letters=None, forbidden_letters=None):
    """
    Trie单词检索
    :param required_letters: 单词必须包含的全部字母，传None表示无强制包含要求
    :param forbidden_letters: 单词不能包含的字母，传None表示无禁用要求
    :return: 符合规则的单词列表
    """
    # 统一转大写集合做匹配，兼容输入大小写
    required_chars = set(c.upper() for c in required_letters) if required_letters else set()
    forbidden_chars = set(c.upper() for c in forbidden_letters) if forbidden_letters else set()
    match_words = []

    def dfs(node, current_word, collected_chars):
        # 到达单词末尾，校验必须包含的字母是否全部存在
        if '' in node and required_chars.issubset(collected_chars):
            match_words.append(current_word)
        # 遍历子节点做递归
        for char, child in node.items():
            if char == '':
                continue
            # 碰到禁用字母直接剪枝，跳过整个分支
            if char in forbidden_chars:
                continue
            dfs(child, current_word + char, collected_chars | {char})
    
    dfs(trie_root, "", set())
    return match_words

使用示例 #

用给出的测试词表验证效果：

test_words = ["APPLE", "EGG", "CAR", "BLUE", "AGRICULTURE", "DONE"]
test_trie = load_trie(test_words)

# 查询包含A、G的单词
print(search_words(test_trie, required_letters=["A", "G"]))
# 输出: ['AGRICULTURE']

# 查询不包含E、G的单词
print(search_words(test_trie, forbidden_letters=["E", "G"]))
# 输出: ['CAR']

# 加载sowpods词表的TRIE后，直接传入TRIE作为第一个参数即可查询全量词典
# 例如查询所有包含X、Z且不包含元音字母的单词
# print(search_words(TRIE, required_letters=["X","Z"], forbidden_letters=["A","E","I","O","U"]))

实现说明 #

• 遍历过程中遇到禁用字母直接剪枝，跳过该分支下所有可能的单词，大词表下性能远高于全量导出后过滤
• 支持同时传入必须包含、禁止包含两类规则，两类条件会同时生效
• 自动对输入的规则字母做大写转换，和Trie中存储的大写单词匹配，不需要调用方手动处理大小写

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