Solution: Trie Implementation
class TrieNode:
"""Node in Trie"""
def __init__(self):
self.children = {} # char -> TrieNode
self.is_end_of_word = False
class Trie:
"""
Trie (Prefix Tree) implementation
Operations:
- insert(word): Add word to trie - O(L)
- search(word): Check if word exists - O(L)
- startsWith(prefix): Check if prefix exists - O(L)
Space: O(ALPHABET_SIZE * N * L) worst case
"""
def __init__(self):
self.root = TrieNode()
def insert(self, word: str) -> None:
"""
Insert word into trie
Time: O(L) where L is length of word
"""
node = self.root
for char in word:
if char not in node.children:
node.children[char] = TrieNode()
node = node.children[char]
node.is_end_of_word = True
def search(self, word: str) -> bool:
"""
Check if word exists in trie (exact match)
Time: O(L)
"""
node = self.root
for char in word:
if char not in node.children:
return False
node = node.children[char]
return node.is_end_of_word
def startsWith(self, prefix: str) -> bool:
"""
Check if any word starts with prefix
Time: O(L)
"""
node = self.root
for char in prefix:
if char not in node.children:
return False
node = node.children[char]
return True
# Test
trie = Trie()
trie.insert("apple")
print(trie.search("apple")) # True
print(trie.search("app")) # False
print(trie.startsWith("app")) # True
trie.insert("app")
print(trie.search("app")) # True
Advanced: Trie with Additional Features
class AdvancedTrie:
"""Trie with word count and autocomplete"""
class Node:
def __init__(self):
self.children = {}
self.is_end = False
self.word_count = 0 # How many times word inserted
def __init__(self):
self.root = self.Node()
def insert(self, word: str) -> None:
"""Insert word and track count"""
node = self.root
for char in word:
if char not in node.children:
node.children[char] = self.Node()
node = node.children[char]
node.is_end = True
node.word_count += 1
def count_words(self, word: str) -> int:
"""Return how many times word was inserted"""
node = self.root
for char in word:
if char not in node.children:
return 0
node = node.children[char]
return node.word_count if node.is_end else 0
def autocomplete(self, prefix: str, limit: int = 10) -> list:
"""
Find all words starting with prefix
Time: O(L + N) where N is number of words with prefix
"""
# Navigate to prefix node
node = self.root
for char in prefix:
if char not in node.children:
return []
node = node.children[char]
# DFS to find all words
results = []
def dfs(node, path):
if len(results) >= limit:
return
if node.is_end:
results.append(prefix + path)
for char, child in node.children.items():
dfs(child, path + char)
dfs(node, "")
return results
def delete(self, word: str) -> bool:
"""
Delete word from trie
Time: O(L)
"""
def _delete(node, word, index):
if index == len(word):
if not node.is_end:
return False
node.is_end = False
return len(node.children) == 0
char = word[index]
if char not in node.children:
return False
should_delete_child = _delete(node.children[char], word, index + 1)
if should_delete_child:
del node.children[char]
return not node.is_end and len(node.children) == 0
return False
return _delete(self.root, word, 0)
# Test autocomplete
trie = AdvancedTrie()
words = ["apple", "app", "application", "apply", "ape", "banana"]
for word in words:
trie.insert(word)
print(trie.autocomplete("app"))
# Output: ['app', 'apple', 'application', 'apply']
print(trie.autocomplete("ap"))
# Output: ['app', 'apple', 'application', 'apply', 'ape']
Optimized: Array-based Trie
class ArrayTrie:
"""
Trie using arrays instead of hash maps
Faster for lowercase letters (26 children)
"""
class Node:
def __init__(self):
self.children = [None] * 26
self.is_end = False
def __init__(self):
self.root = self.Node()
def _char_to_index(self, char):
return ord(char) - ord('a')
def insert(self, word: str) -> None:
node = self.root
for char in word:
idx = self._char_to_index(char)
if not node.children[idx]:
node.children[idx] = self.Node()
node = node.children[idx]
node.is_end = True
def search(self, word: str) -> bool:
node = self.root
for char in word:
idx = self._char_to_index(char)
if not node.children[idx]:
return False
node = node.children[idx]
return node.is_end
Complexity Analysis
- Insert: O(L) - traverse L nodes
- Search: O(L) - traverse L nodes
- StartsWith: O(L) - traverse L nodes
- Space: O(ALPHABET * N * L) worst case
- Hash map: More space-efficient for sparse tries
- Array: Faster access but wastes space if sparse
Common Interview Follow-ups
Q: How to handle case-insensitive search?
A: Convert to lowercase on insert/search, or increase alphabet size to 52.
Q: How to implement delete?
A: Recursive approach removing nodes that become unused.
Q: How to find longest word with prefix?
A: DFS from prefix node, track longest path to end node.
Q: Memory optimization?
A: Use hash map for sparse tries, compress single-child chains.
Related Problems
