Implementing a Trie to support autocomplete in Python
You could just implement a generator that iterates over the Trie according to prefix the same way as other methods do. Once you've found the node at the end of the prefix you can use recursive generator with yield from
to iterate over the sub trie while keeping track of the prefix and yielding it when terminal node is found:
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')))
Output:
['foo', 'foob', 'foobar', 'foof']
from collections import defaultdict
class TrieNode:
def __init__(self):
self.node = defaultdict(TrieNode)
self.is_word = False
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, words):
curr = self.root
for char in words:
curr = curr.node[char]
curr.is_word = True
def search(self, word):
curr = self.root
for char in word:
if char not in curr.node:
return False
curr = curr.node[char]
return curr.is_word
def dfs(self, node, word, word_list):
if node.is_word == True:
word_list.append(word)
for a, n in node.node.items():
self.dfs(n, word + a, word_list)
def auto_complete(self, word_to_search, word_list):
temp_word = ""
curr = self.root
for char in word_to_search:
if char not in curr.node.keys():
print("Invalid Input")
else:
temp_word += char
curr = curr.node[char]
self.dfs(curr, temp_word, word_list)
print(word_list)