Insert
The
trie is a kind of tree ideal for implementing autocomplete.
This is how we can
populate our trie with data.
class TrieNode:
def __init__(self):
self.children = {}
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
currentNode = self.root
for c in word:
if currentNode.children.get(c):
currentNode = currentNode.children[c]
else:
newNode = TrieNode()
currentNode.children[c] = newNode
currentNode = newNode
T = Trie()
T.insert('ace')
T.insert('bad')
T.insert('act')
print(T.root.children)
print("a", T.root.children['a'].children)
print("a:c", T.root.children['a'].children['c'].children)
Search
We check if the current node has any children with the
current character as the key.
If there is such a child, we
update the current node to be the child node.
class TrieNode:
def __init__(self):
self.children = {}
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
currentNode = self.root
for c in word:
if currentNode.children.get(c):
currentNode = currentNode.children[c]
else:
newNode = TrieNode()
currentNode.children[c] = newNode
currentNode = newNode
currentNode.children["*"] = None
def search(self, word):
currentNode = self.root
for c in word:
if currentNode.children.get(c):
currentNode = currentNode.children.get(c)
else:
return None
return currentNode
T = Trie()
T.insert('ace')
T.insert('bad')
T.insert('act')
T.insert('cat')
T.insert('bat')
T.insert('batter')
print(T.search('cat').children.items())
print(T.search('bat').children.items())
Autocomplete
We can implement our autocomplete
feature.
class TrieNode:
def __init__(self):
self.children = {}
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
currentNode = self.root
for c in word:
if currentNode.children.get(c):
currentNode = currentNode.children[c]
else:
newNode = TrieNode()
currentNode.children[c] = newNode
currentNode = newNode
currentNode.children["*"] = None
def search(self, word):
currentNode = self.root
for char in word:
if currentNode.children.get(char):
currentNode = currentNode.children.get(char)
else:
return None
return currentNode
def collectWords(self, node=None, prefix=, words=[]):
currentNode = node or self.root
for key, node in currentNode.children.items():
if key == '*':
words.append(prefix)
else:
self.collectWords(node, prefix + key, words)
return words
def autocomplete(self, prefix):
node = self.search(prefix)
if not node:
return None
return self.collectWords(node, prefix, words=[])
T = Trie()
for word in ['ace', 'bad', 'act', 'cat', 'bat', 'batter']:
T.insert(word)
words = T.autocomplete('bat')
print("Autocomplete words for 'bat' = ", words)
words = T.autocomplete('ac')
print("Autocomplete words for 'ac' = ", words)