Prefix Search and Word Problems
28 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Once words sit in a trie, a prefix points straight to the part of the tree that holds the answer. This lesson lists, counts and shortens words by prefix, searches with a wildcard, and shows where bisect is the simpler tool.
All words under a prefix
Depth-first search follows one branch as far as it goes before trying the next, usually with recursion. In a trie it visits every node below a starting point once.
Autocomplete lists the words that start with what a user typed. Walk down the prefix to its node. Then do a depth-first search below it, adding a letter at each step and collecting the nodes marked as word ends.
def words_with_prefix(trie, prefix): result = []
def dfs(node, path): if node.is_end: result.append(path) for ch in sorted(node.children): dfs(node.children[ch], path + ch)
start = trie._walk(prefix) if start: dfs(start, prefix) return resultExample 1 · Autocomplete
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 _walk(self, text): node = self.root for ch in text: node = node.children.get(ch) if node is None: return None return node
def search(self, word): node = self._walk(word) return node is not None and node.is_end
def starts_with(self, prefix): return self._walk(prefix) is not None
def words_with_prefix(trie, prefix): result = []
def dfs(node, path): if node.is_end: result.append(path) for ch in sorted(node.children): dfs(node.children[ch], path + ch)
start = trie._walk(prefix) if start: dfs(start, prefix) return result
trie = Trie()for word in ["cart", "car", "cat", "dog", "cattle"]: trie.insert(word)print(words_with_prefix(trie, "ca"))print(words_with_prefix(trie, "cat"))print(words_with_prefix(trie, "x"))['car', 'cart', 'cat', 'cattle'] ['cat', 'cattle'] []
What this code does
- "car" comes before "cart" because a node is added to the result before its children are visited.
- A prefix that is not in the trie gives an empty list, and "cat" is included because it is itself a stored word.
Example 2 · Only the first three suggestions
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 _walk(self, text): node = self.root for ch in text: node = node.children.get(ch) if node is None: return None return node
def search(self, word): node = self._walk(word) return node is not None and node.is_end
def starts_with(self, prefix): return self._walk(prefix) is not None
def suggest(trie, prefix, limit): result = []
def dfs(node, path): if len(result) == limit: return if node.is_end: result.append(path) for ch in sorted(node.children): dfs(node.children[ch], path + ch)
start = trie._walk(prefix) if start: dfs(start, prefix) return result
trie = Trie()for word in ["tea", "ten", "tent", "tend", "tens", "to"]: trie.insert(word)print(suggest(trie, "te", 3))['tea', 'ten', 'tend']
What this code does
- The walk stops early once three words are found, so the cost is not tied to how many words there are below.
- A guard at the top of the function is simpler than checking the size inside the loop.
Visiting children in sorted letter order gives the words in alphabetical order. The cost is O(L) for the walk plus the size of the part of the trie below the prefix.
Counting words with a prefix
To count the words with a prefix without listing them, store a count in each node: how many words pass through it. Insert adds one to every node on its path. The answer is the count at the prefix node.
class CNode: def __init__(self): self.children = {} self.count = 0 self.is_end = False
def insert(root, word): node = root for ch in word: node = node.children.setdefault(ch, CNode()) node.count += 1 node.is_end = TrueExample 3 · Count words with a prefix
class CNode: def __init__(self): self.children = {} self.count = 0 self.is_end = False
def insert(root, word): node = root for ch in word: node = node.children.setdefault(ch, CNode()) node.count += 1 node.is_end = True
def count_prefix(root, prefix): node = root for ch in prefix: node = node.children.get(ch) if node is None: return 0 return node.count
root = CNode()for word in ["car", "cart", "cat", "dog"]: insert(root, word)for prefix in ["ca", "car", "d", "z"]: print(prefix, count_prefix(root, prefix))ca 3 car 2 d 1 z 0
What this code does
- The count at the node for "ca" is 3, as three words pass through it. A missing prefix returns 0.
- Each query costs O(L) and does not visit the nodes below.
Insert is still O(L). Reading the count is O(L) too, and the answer is ready without visiting the nodes below. Duplicate words would be counted twice, so decide first whether the input has duplicates.
Shortest unique prefix
Each name needs the shortest start that tells it apart. Store a count in every node and walk the name from the root. The first node with a count of 1 is used by no other name, so stop there.
def shortest_prefix(root, word): node = root for i, ch in enumerate(word): node = node.children[ch] if node.count == 1: return word[:i + 1] return wordExample 4 · Shortest unique prefixes
class CNode: def __init__(self): self.children = {} self.count = 0 self.is_end = False
def insert(root, word): node = root for ch in word: node = node.children.setdefault(ch, CNode()) node.count += 1 node.is_end = True
def count_prefix(root, prefix): node = root for ch in prefix: node = node.children.get(ch) if node is None: return 0 return node.count
def shortest_prefix(root, word): node = root for i, ch in enumerate(word): node = node.children[ch] if node.count == 1: return word[:i + 1] return word
names = ["zebra", "dog", "duck", "dove"]root = CNode()for name in names: insert(root, name)for name in names: print(name, shortest_prefix(root, name))zebra z dog dog duck du dove dov
What this code does
- "zebra" is the only word starting with z. "dog", "duck" and "dove" share d, so they need two or three letters.
- "do" is shared by dog and dove, so those two need three letters.
This assumes every word is different, and that no word is a prefix of another. If one is, it returns the whole word. The cost is O(L) per word.
Wildcards, and bisect as another tool
A search such as c.t lets a dot match any one letter. At a dot, try every child. At a letter, follow that child only. The search succeeds if any branch reaches the end of the pattern on a word end.
def match(node, pattern, i=0): if i == len(pattern): return node.is_end ch = pattern[i] if ch == ".": return any(match(c, pattern, i + 1) for c in node.children.values()) child = node.children.get(ch) return child is not None and match(child, pattern, i + 1)Example 5 · Wildcard match
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 _walk(self, text): node = self.root for ch in text: node = node.children.get(ch) if node is None: return None return node
def search(self, word): node = self._walk(word) return node is not None and node.is_end
def starts_with(self, prefix): return self._walk(prefix) is not None
def match(node, pattern, i=0): if i == len(pattern): return node.is_end ch = pattern[i] if ch == ".": return any(match(c, pattern, i + 1) for c in node.children.values()) child = node.children.get(ch) return child is not None and match(child, pattern, i + 1)
trie = Trie()for word in ["bad", "dad", "mad"]: trie.insert(word)for pattern in ["pad", "bad", ".ad", "b..", "b.", "...."]: print(pattern, match(trie.root, pattern))pad False bad True .ad True b.. True b. False .... False
What this code does
- ".ad" matches all three words, and "b.." matches bad. "b." is too short and "...." is too long, so both fail.
- A pattern longer than every word runs out of nodes, so the answer is False.
Example 6 · Prefix search with bisect
import bisect
words = sorted(["cart", "car", "cat", "dog", "cattle", "ant"])prefix = "ca"start = bisect.bisect_left(words, prefix)end = bisect.bisect_left(words, prefix + chr(0x10FFFF))print(words[start:end], end - start)['car', 'cart', 'cat', 'cattle'] 4
What this code does
- Words with the same start sit side by side in a sorted list. Two binary searches find the block.
- This costs O(log n) per query plus the output. It is shorter than a trie, but it needs the list sorted first and does not allow cheap inserts.
A dot can branch into many nodes, so the worst case is the whole trie. For plain prefix questions on a fixed list, a sorted list with bisect is shorter and uses less memory. Interviews still expect the trie.
Next step
Check your understanding
10 quick questions