Prefix Tree Operations
Problem
Build a prefix tree for lowercase words with three operations. insert(word) stores a word. has_word(word) reports whether exactly that word was stored. has_prefix(prefix) reports whether any stored word starts with the given non-empty prefix. Each operation should cost time proportional to the length of its argument, however many words are stored.
Examples
Input: insert("apple"), has_word("apple"), has_word("app"), has_prefix("app")
Output: True, False, True
Why: "app" begins a stored word but was never stored itselfInput: insert("apple"), insert("app"), has_word("app")
Output: TrueInput: (nothing inserted), has_word("app"), has_prefix("a")
Output: False, False
Why: edge case, an empty tree contains no words and no prefixesHints
0 / 3
Checking a prefix against every stored word costs time for every word. Words that begin the same way should share the work of storing that beginning.
Store the words as paths of nested children, one character per step, so that words with a common prefix share the same first nodes. A path existing is not the same as a word ending there, so the end of a word needs its own marker.
Insert by walking from the root and creating each missing child, then set an end marker on the last node. For both queries walk the same way and report failure as soon as a character has no child. A prefix query succeeds if the walk finishes; a word query also needs the end marker on the final node.
Solution
Each node is a dictionary from a character to the child node, so a word is a path from the root and shared prefixes share nodes. A sentinel key marks the nodes where a stored word actually ends, which is what separates a stored "app" from a mere prefix of "apple". Both queries use the same walk and differ only in whether they demand the sentinel at the end. Every operation touches one node per character, so each costs O(L) for an argument of length L, and the tree uses O(total characters stored) space.
class PrefixTree:
def __init__(self):
self.root = {}
def insert(self, word):
node = self.root
for ch in word:
node = node.setdefault(ch, {}) # shared prefixes share nodes
node["$"] = True # a whole word ends here
def _walk(self, s): # the node s leads to, or None
node = self.root
for ch in s:
if ch not in node:
return None
node = node[ch]
return node
def has_word(self, word):
node = self._walk(word)
return node is not None and "$" in node
def has_prefix(self, prefix):
return self._walk(prefix) is not None
t = PrefixTree(); t.insert("apple")
print(t.has_word("apple"), t.has_word("app"), t.has_prefix("app")) # -> True False True
t.insert("app"); e = PrefixTree()
print(t.has_word("app"), e.has_word("app"), e.has_prefix("a")) # -> True False FalseStuck on the idea rather than the code? Trie Basics covers it.