2017-09-04 35 views
2

我試圖實現支持網站上自動完成的數據結構。 我設法實現了一個Trie的迭代版本。它支持在Trie中添加和搜索的兩種主要方法。 但是現在我需要添加一個方法,返回以下列前綴開頭的所有單詞。有人可以幫我弄這個嗎。實現Trie以支持Python中的自動完成

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. 

回答

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']