DSA course menu
Tries (Prefix Trees) in Python: Autocomplete and Word Search
Learn the trie data structure: storing words character by character, insert and search in O(word length), prefix searches for autocomplete, counting words with a prefix, and when to use a trie instead of a set.
On this page
A trie (pronounced "try", from retrieval), also called a prefix tree, stores a set of words so that words sharing a beginning share the same path. Each edge is a character; walking from the root spells out a prefix.
(root)
/ \
c d
| |
a o
/ \ |
t* r* g*
|
s*This trie holds "cat", "car", "cars" and "dog". A * marks nodes where a complete word ends: "car" is a word, and it is also the prefix of "cars".
Building a trie#
Each node keeps a dictionary from character to child node, plus a flag for "a word ends here":
class TrieNode:
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 _walk(self, text):
node = self.root
for ch in text:
node = node.children.get(ch)
if node is None:
return None
return node
def contains(self, word):
node = self._walk(word)
return node is not None and node.is_word
def starts_with(self, prefix):
return self._walk(prefix) is not None
t = Trie()
for w in ["cat", "car", "cars", "dog"]:
t.insert(w)
print(t.contains("car"), t.contains("ca"), t.contains("cow"))
print(t.starts_with("ca"), t.starts_with("do"), t.starts_with("x"))True False False
True True False"ca" is a valid prefix but not a stored word, which is exactly the difference the is_word flag captures.
Insert and search cost O(L) for a word of length L, no matter how many words are stored.
Autocomplete#
To list every word starting with a prefix, walk to the prefix's node, then collect all words below it with a depth-first search:
class TrieNode:
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 complete(self, prefix):
node = self.root
for ch in prefix:
if ch not in node.children:
return []
node = node.children[ch]
results = []
def collect(n, path):
if n.is_word:
results.append(prefix + path)
for ch in sorted(n.children):
collect(n.children[ch], path + ch)
collect(node, "")
return results
t = Trie()
for w in ["python", "pythonic", "pyramid", "pandas", "panda", "java", "javascript"]:
t.insert(w)
print(t.complete("py"))
print(t.complete("pand"))
print(t.complete("java"))
print(t.complete("ruby"))['pyramid', 'python', 'pythonic']
['panda', 'pandas']
['java', 'javascript']
[]Visiting children in sorted order returns the suggestions alphabetically.
Counting words with a prefix#
Store, on each node, how many words pass through it. Then "how many words start with this prefix?" is O(prefix length):
class TrieNode:
def __init__(self):
self.children = {}
self.passing = 0
class CountingTrie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
node = self.root
for ch in word:
node = node.children.setdefault(ch, TrieNode())
node.passing += 1
def count_prefix(self, prefix):
node = self.root
for ch in prefix:
node = node.children.get(ch)
if node is None:
return 0
return node.passing
t = CountingTrie()
for w in ["apple", "app", "application", "apply", "banana"]:
t.insert(w)
print(t.count_prefix("app"), t.count_prefix("appl"), t.count_prefix("b"), t.count_prefix("c"))4 3 1 0Trie or hash set?#
Use a trie when prefixes matter: autocomplete, spell checkers, IP routing tables, word games. For plain membership, a set is simpler and smaller.
Common mistakes#
Practice#
Practice exercise
Write longest_common_prefix(words) using a trie: insert all words, then walk down from the root while the current node has exactly one child and is not the end of a word.
Show hint
Build the trie with an is_word flag. Stop as soon as a node branches (more than one child) or a word ends there, because the common prefix cannot be longer than any word.
Show solution
class TrieNode:
def __init__(self):
self.children = {}
self.is_word = False
def longest_common_prefix(words):
if not words:
return ""
root = TrieNode()
for w in words:
node = root
for ch in w:
node = node.children.setdefault(ch, TrieNode())
node.is_word = True
prefix, node = "", root
while len(node.children) == 1 and not node.is_word:
ch, node = next(iter(node.children.items()))
prefix += ch
return prefix
print(repr(longest_common_prefix(["flower", "flow", "flight"])))
print(repr(longest_common_prefix(["interview", "internet", "interval"])))
print(repr(longest_common_prefix(["dog", "racecar"])))
print(repr(longest_common_prefix(["same", "same"])))'fl'
'inter'
''
'same'Key takeaway#
Next: networks of connections, in Graphs and Their Representation.
Finished this lesson?