实现Trie以支持Python中的自动完成

Question

我试图实现支持网站上自动完成的数据结构。 我设法实现了一个Trie的迭代版本。它支持在Trie中添加和搜索的两种主要方法。 但是现在我需要添加一个方法，返回以下列前缀开头的所有单词。有人可以帮我弄这个吗。

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

    def insert(self, word): 
     curr = self.root 
     for letter in word: 
      node = curr.children.get(letter) 
      if not node: 
       node = TrieNode() 
       curr.children[letter] = node 
      curr = node 
     curr.end = True 

    def search(self, word): 
     curr = self.root 
     for letter in word: 
      node = curr.children.get(letter) 
      if not node: 
       return False 
      curr = node 
     return curr.end 

    def all_words_beginning_with_prefix(self, prefix): 
     #I'm not sure how to go about this one. 

Answer 1

你可以实现一个生成器，该生成器根据前缀按照与其他方法相同的方式迭代Trie。一旦你找到节点的前缀结束时，您可以使用递归发生器yield from遍历子特里同时跟踪前缀和产生它时，终端节点发现：

class TrieNode: 
    def __init__(self): 
     self.end = False 
     self.children = {} 

    def all_words(self, prefix): 
     if self.end: 
      yield prefix 

     for letter, child in self.children.items(): 
      yield from child.all_words(prefix + letter) 

class Trie: 
    # existing methods here 
    def all_words_beginning_with_prefix(self, prefix): 
     cur = self.root 
     for c in prefix: 
      cur = cur.children.get(c) 
      if cur is None: 
       return # No words with given prefix 

     yield from cur.all_words(prefix) 

trie = Trie() 
trie.insert('foobar') 
trie.insert('foo') 
trie.insert('bar') 
trie.insert('foob') 
trie.insert('foof') 

print(list(trie.all_words_beginning_with_prefix('foo'))) 

输出：

['foo', 'foob', 'foobar', 'foof']