如何在Python的Trie字典中返回单词对应频率值

问题背景 #

我需要在Trie数据结构中搜索指定单词，若单词存在则返回该单词及其对应的频率。当前打印self.root.children得到的Trie结构如下：

{'c': {'u': {'t': {('e', 10): True, ('s', 50): True},
   ('t', 30): True,
   ('b', 15): True},
  'a': {'l': {('m', 1000): True}}},
 'a': {'n': {('t', 20): True,
   'n': {'o': {'t': {'a': {'t': {'i': {'o': {('n', 5): True}}}}}}}},
  'p': {'p': {'l': {('e', 300): True},
    'e': {'n': {'d': {'i': {('x', 10): True}}}},
    ('s', 60): True},
   'o': {'l': {'o': {'g': {('y', 600): True,
       'e': {'t': {'i': {('c', 1000): True}}}}}}}},
 'f': {'u': {'r': {'t': {'h': {'e': {('r', 40): True}}},
    'n': {'i': {'t': {'u': {'r': {('e', 500): True}}}}}},
   ('r', 10): True},
  'i': {'n': {('d', 400): True}},
  'a': {'r': {('m', 5000): True,
    'm': {'i': {'n': {('g', 1000): True}}, 'e': {('r', 300): True}}},
   't': {'h': {'o': {('m', 40): True}}}}}}

例如，搜索单词cuts时，需要返回对应的数值50（注：True标记单词的结束位置）。

我目前写的搜索函数伪代码如下：

def search(self, word: str) -> int:

    current = self.root

    if word in current.children:
        return its integer pair if its the end of a word and matches
            
    return 0

Trie类的定义如下：

# Class representing a node in the Trie
class TrieNode:
    def __init__(self, letter=None, frequency=None, is_last=False):
        self.letter = letter            # letter stored at this node
        self.frequency = frequency      # frequency of the word if this letter is the end of a word
        self.is_last = is_last          # True if this letter is the end of a word
        self.children: dict[str, TrieNode] = {}     # a hashtable containing children nodes, key = letter, value = child node

class TrieDictionary(BaseDictionary):
    def __init__(self):
        # setting root to object TrieNode()
        self.root = TrieNode()

实现方案 #

要实现正确的搜索逻辑，需要逐个遍历单词的每个字符，沿着Trie的节点层级向下查找，最后判断当前节点是否为单词的结尾并返回对应频率。具体代码实现如下：

def search(self, word: str) -> int:
    current = self.root
    for char in word:
        # 如果当前字符不在子节点中，说明单词不存在，返回0
        if char not in current.children:
            return 0
        # 移动到下一个节点
        current = current.children[char]
    # 遍历完所有字符后，检查当前节点是否是单词结尾
    if current.is_last:
        return current.frequency
    # 不是结尾则返回0
    return 0

代码说明 #

1. 从根节点开始，遍历目标单词的每一个字符：
   • 若某个字符不在当前节点的children字典中，直接返回0（单词不存在）
   • 否则移动到该字符对应的子节点，继续下一个字符的遍历
2. 遍历完所有字符后，检查当前节点的is_last标记：
   • 若为True，说明这是一个完整单词的结尾，返回节点的frequency值
   • 若为False，说明该字符序列只是某个长单词的前缀，不是目标单词，返回0

比如搜索cuts时，会依次遍历c→u→t→s，最后到达的节点is_last为True，对应的frequency是50，因此返回50。

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