Trie Implementation in Python
A trie (prefix tree) stores strings character-by-character so prefix queries run in O(length) instead of O(N * length). The Pythonic shape uses a dict of dicts plus a sentinel key for 'word ends here'. This entry covers the dict-of-dicts trie, a class-based variant with explicit nodes, and the autocomplete query you build it for.
697 views
21
END = '__end__'
def build_trie(words):
root = {}
for word in words:
node = root
for ch in word:
node = node.setdefault(ch, {})
node[END] = True
return root
def contains(trie, word):
node = trie
for ch in word:
if ch not in node:
return False
node = node[ch]
return node.get(END, False)
def has_prefix(trie, prefix):
node = trie
for ch in prefix:
if ch not in node:
return False
node = node[ch]
return True
t = build_trie(['cat', 'car', 'card', 'care', 'dog'])
print(contains(t, 'car')) # True
print(contains(t, 'card')) # True
print(contains(t, 'ca')) # False (prefix only, not a stored word)
print(contains(t, 'cards')) # False (no word continues past 'card')
print(has_prefix(t, 'ca')) # True
print(has_prefix(t, 'do')) # True
print(has_prefix(t, 'fish')) # FalseThe dict-of-dicts shape is the simplest trie that works: every node is a dict, keys are next characters, values are child dicts. A sentinel key ('__end__' here) marks 'a word ends at this node', which is what distinguishes a stored word from a mere prefix. setdefault is the one-liner for 'get or create the next node'. Insertion and lookup are both O(length) and the storage cost is one dict entry per character of input. This is the right starting trie for autocomplete prototypes, spell checkers, and the LeetCode 'Implement Trie' family.
class TrieNode:
__slots__ = ('children', 'is_word')
def __init__(self):
self.children = {}
self.is_word = 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_word = True
def search(self, word):
node = self._descend(word)
return node is not None and node.is_word
def starts_with(self, prefix):
return self._descend(prefix) is not None
def _descend(self, s):
node = self.root
for ch in s:
node = node.children.get(ch)
if node is None:
return None
return node
t = Trie()
for word in ['apple', 'app', 'apply']:
t.insert(word)
print(t.search('app')) # True
print(t.search('apple')) # True
print(t.search('apples')) # False
print(t.starts_with('appl')) # True
print(t.starts_with('banana')) # FalseThe class-based form swaps the sentinel-key trick for an explicit is_word boolean per node, which reads more clearly and survives serialization libraries that mangle dict keys with leading underscores. __slots__ cuts the per-node memory footprint in half, which matters for million-word dictionaries (English has roughly 200K words; biology has millions of taxon names). The _descend helper is the shared walk used by both search and starts_with; deduplicating it is what keeps the public methods two lines each.
class TrieNode:
__slots__ = ('children', 'is_word')
def __init__(self):
self.children = {}
self.is_word = False
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
node = self.root
for ch in word:
node = node.children.setdefault(ch, TrieNode())
node.is_word = True
def autocomplete(self, prefix, limit=10):
node = self.root
for ch in prefix:
node = node.children.get(ch)
if node is None:
return []
results = []
# DFS the subtree, collecting words in lexicographic order.
def walk(curr, path):
if len(results) >= limit:
return
if curr.is_word:
results.append(prefix + ''.join(path))
for ch in sorted(curr.children):
if len(results) >= limit:
return
path.append(ch)
walk(curr.children[ch], path)
path.pop()
walk(node, [])
return results
t = Trie()
for word in ['cat', 'car', 'card', 'care', 'cargo', 'carrot', 'dog']:
t.insert(word)
print(t.autocomplete('car')) # ['car', 'card', 'care', 'cargo', 'carrot']
print(t.autocomplete('ca', limit=3)) # ['car', 'card', 'care']
print(t.autocomplete('zzz')) # [] (no words with that prefix)
print(t.autocomplete('')) # all words, capped at the limitAutocomplete is the trie's killer feature: descend to the prefix node, then DFS the subtree to enumerate every stored word under it. Sorting children at each step gives lexicographic order; iterating node.children directly gives insertion order. The limit parameter and the early-exit checks inside walk keep the function bounded for huge dictionaries, since you almost never want all matches in a UI. Real-world autocomplete adds frequency weights to each node, then uses a heap during the walk to return the K most-frequent suggestions instead of the alphabetically-first K.