Trie enables O(L) prefix search — insert, search, starts_with, and autocomplete in O(word length).
Trie (Prefix Tree)
class TrieNode:
def __init__(self):
self.children = {}
self.is_end = False
class Trie:
def __init__(self): self.root = TrieNode()
def insert(self, word):
node = self.root
for ch in word:
if ch not in node.children:
node.children[ch] = TrieNode()
node = node.children[ch]
node.is_end = True
def search(self, word) -> bool:
node = self.root
for ch in word:
if ch not in node.children: return False
node = node.children[ch]
return node.is_end
def starts_with(self, prefix) -> bool:
node = self.root
for ch in prefix:
if ch not in node.children: return False
node = node.children[ch]
return True
def autocomplete(self, prefix) -> list:
node = self.root
for ch in prefix:
if ch not in node.children: return []
node = node.children[ch]
# DFS from here to find all words
results = []
def dfs(n, path):
if n.is_end: results.append(prefix+path)
for c,child in n.children.items(): dfs(child, path+c)
dfs(node, '')
return results