基於Trie樹實現自動補齊(autocomplete)功能
阿新 • • 發佈:2020-11-21
""" Trie樹 """ __author__ = 'youngf' class TrieNode: def __init__(self): self.children = {} self.last = False class Trie: def __init__(self): self.root = TrieNode() def from_keys(self, words): """ 批量新增詞 """ for word in words: self.insert(word) def insert(self, word): """ 單條新增 """ node = self.root for a in list(word): if not node.children.get(a): node.children[a] = TrieNode() node = node.children[a] node.last = True def search_word(self, word): """ 全匹配搜尋 """ node = self.root if not node: return False for a in list(word): if not node.children.get(a): return False node = node.children[a] return node.last def search_prefix(self, prefix): """ 字首搜尋 """ node = self.root if not node: return False for a in list(prefix): if not node.children.get(a): return False node = node.children[a] return True def auto_complete(self, prefix): """ 返回所有以prefix為字首的詞 """ def iter_words(node, word, word_list): if node.last: word_list.append(word) # trick: 只要不改變word_list的id,就能保持引用傳遞 if node.children: for ch, children_node in node.children.items(): # 深度優先遍歷 iter_words(children_node, word + ch, word_list) node = self.root temp_word = '' candidate_words = [] for ch in list(prefix): if not node.children.get(ch): # raise Exception('no such prefix') return [] temp_word += ch node = node.children[ch] iter_words(node, temp_word, candidate_words) return candidate_words if __name__ == '__main__': t = Trie() t.from_keys(['her', 'hell', 'hello', 'hers', 'cat']) print(t.search_prefix('he')) print(t.search_word('he')) print(t.auto_complete('he'))