1"""
2Trie (Prefix Tree) - String-Optimized Tree Data Structure
3
4A Trie is a tree-like data structure used to store strings where each node
5represents a character. The path from root to any node represents a prefix
6of the stored strings. Nodes can be marked as "end of word" to indicate
7complete words.
8
9Key Properties:
10 1. Root node is empty (represents empty string prefix).
11 2. Each node has up to 26 children (for lowercase a-z) or more for extended alphabets.
12 3. Nodes marked as end_of_word indicate complete words.
13 4. Common prefixes share the same path from root.
14
15Time Complexity (where m = length of word/prefix):
16 - insert: O(m)
17 - search: O(m)
18 - delete: O(m)
19 - starts_with: O(m)
20 - get_all_words: O(n * m) where n = number of words
21
22Space Complexity:
23 - Storage: O(ALPHABET_SIZE * m * n) worst case, but typically much less due to prefix sharing
24 - Operations: O(m) for insert/search/delete, O(n * m) for get_all_words
25"""
26
27from enum import Enum
28
29
30class TrieNode:
31 """A single node in the trie."""
32
33 __slots__ = ("children", "end_of_word", "char")
34
35 def __init__(self, char: str = ""):
36 """Initialize a trie node with a character."""
37 self.char = char # Character stored at this node (empty for root)
38 self.children: dict[str, "TrieNode"] = {} # Map from char to child node
39 self.end_of_word: bool = False # True if this node marks end of a word
40
41 def __str__(self):
42 return f"TrieNode('{self.char}', end={self.end_of_word})"
43
44 def __repr__(self):
45 return self.__str__()
46
47
48class Trie:
49 """
50 Trie (Prefix Tree) implementation.
51
52 Provides operations for inserting, searching, and deleting strings,
53 as well as prefix matching and traversal methods.
54 """
55
56 def __init__(self):
57 """Initialize an empty trie with a root node."""
58 self.root = TrieNode() # Root has empty character
59
60 def __contains__(self, word: str) -> bool:
61 """Allow 'word in trie' syntax."""
62 return self.search(word)
63
64 # -------------------------------------------------------------------------
65 # Core Operations
66 # -------------------------------------------------------------------------
67
68 def insert(self, word: str) -> None:
69 """
70 Insert a word into the trie.
71 Creates new nodes as needed along the path.
72 """
73 if not word:
74 return
75
76 node = self.root
77 for char in word:
78 if char not in node.children:
79 node.children[char] = TrieNode(char)
80 node = node.children[char]
81 node.end_of_word = True
82
83 def search(self, word: str) -> bool:
84 """
85 Search for a complete word in the trie.
86 Returns True if the word exists, False otherwise.
87 """
88 node = self._find_node(word)
89 return node is not None and node.end_of_word
90
91 def delete(self, word: str) -> bool:
92 """
93 Delete a word from the trie.
94 Returns True if deletion was successful, False if word not found.
95
96 Handles three cases:
97 1. Word not found - return False
98 2. Word is prefix of another word - only unmark end_of_word
99 3. Word has unique suffix - remove nodes that are not part of other words
100 """
101 if not word:
102 return False
103
104 def _delete_recursive(node: TrieNode, word: str, depth: int) -> bool:
105 """
106 Recursively delete word from trie.
107 Returns True if current node should be deleted (has no children and not end of word).
108 """
109 if depth == len(word):
110 # Reached end of word
111 if not node.end_of_word:
112 return False # Word doesn't exist
113 node.end_of_word = False
114 # Delete node if it has no children
115 return len(node.children) == 0
116
117 char = word[depth]
118 if char not in node.children:
119 return False # Word doesn't exist
120
121 child = node.children[char]
122 should_delete_child = _delete_recursive(child, word, depth + 1)
123
124 if should_delete_child:
125 del node.children[char]
126 # Current node can be deleted if it's not end of word and has no children
127 return not node.end_of_word and len(node.children) == 0
128
129 return False
130
131 # Start deletion from root
132 _delete_recursive(self.root, word, 0)
133 return True
134
135 def starts_with(self, prefix: str) -> bool:
136 """
137 Check if any word in the trie starts with the given prefix.
138 Returns True if prefix exists, False otherwise.
139 """
140 return self._find_node(prefix) is not None
141
142 def _find_node(self, prefix: str) -> TrieNode | None:
143 """
144 Find the node corresponding to the last character of prefix.
145 Returns None if prefix doesn't exist in trie.
146 """
147 if prefix == "":
148 return self.root
149
150 node = self.root
151 for char in prefix:
152 if char not in node.children:
153 return None
154 node = node.children[char]
155 return node
156
157 # -------------------------------------------------------------------------
158 # Word Retrieval Methods
159 # -------------------------------------------------------------------------
160
161 def get_all_words(self) -> list[str]:
162 """
163 Get all complete words stored in the trie.
164 Returns a list of all words in lexicographical order.
165 """
166 words = []
167 self._collect_words(self.root, "", words)
168 return words
169
170 def get_words_with_prefix(self, prefix: str) -> list[str]:
171 """
172 Get all words that start with the given prefix.
173 Returns empty list if prefix doesn't exist.
174 """
175 node = self._find_node(prefix)
176 if node is None:
177 return []
178
179 words = []
180 self._collect_words(node, prefix, words)
181 return words
182
183 def _collect_words(self, node: TrieNode, prefix: str, words: list[str]) -> None:
184 """Recursively collect all words starting from node with given prefix."""
185 if node.end_of_word:
186 words.append(prefix)
187
188 # Process children in sorted order for consistent output
189 for char in sorted(node.children.keys()):
190 self._collect_words(node.children[char], prefix + char, words)
191
192 # -------------------------------------------------------------------------
193 # Traversal Methods
194 # -------------------------------------------------------------------------
195
196 class TraverseOrder(Enum):
197 """Enumeration of trie traversal orders."""
198
199 preorder = 1 # Visit node, then children (depth-first)
200 postorder = 2 # Visit children, then node (depth-first)
201 levelorder = 3 # Breadth-first traversal
202
203 def traverse(self, order: TraverseOrder) -> list[str]:
204 """
205 Traverse the trie in the specified order.
206 Returns a list of characters visited (including empty root marker).
207 """
208 result = []
209
210 def _preorder(node: TrieNode):
211 result.append(node.char if node.char else "[root]")
212 for char in sorted(node.children.keys()):
213 _preorder(node.children[char])
214
215 def _postorder(node: TrieNode):
216 for char in sorted(node.children.keys()):
217 _postorder(node.children[char])
218 result.append(node.char if node.char else "[root]")
219
220 def _levelorder():
221 if not self.root.children:
222 result.append("[root]")
223 return
224
225 queue = [self.root]
226 while queue:
227 node = queue.pop(0)
228 result.append(node.char if node.char else "[root]")
229 for char in sorted(node.children.keys()):
230 queue.append(node.children[char])
231
232 traversals = {
233 Trie.TraverseOrder.preorder: _preorder,
234 Trie.TraverseOrder.postorder: _postorder,
235 Trie.TraverseOrder.levelorder: _levelorder,
236 }
237
238 if order not in traversals:
239 raise TypeError("Invalid traverse order passed.")
240
241 if order == Trie.TraverseOrder.levelorder:
242 traversals[order]()
243 else:
244 traversals[order](self.root)
245
246 return result
247
248 # -------------------------------------------------------------------------
249 # Utility Methods
250 # -------------------------------------------------------------------------
251
252 def count_words(self) -> int:
253 """Return the total number of complete words in the trie."""
254 return len(self.get_all_words())
255
256 def count_nodes(self) -> int:
257 """Return the total number of nodes in the trie (including root)."""
258 count = 0
259
260 def _count(node: TrieNode):
261 nonlocal count
262 count += 1
263 for child in node.children.values():
264 _count(child)
265
266 _count(self.root)
267 return count
268
269 def is_empty(self) -> bool:
270 """Return True if the trie contains no words."""
271 return len(self.root.children) == 0
272
273 def longest_common_prefix(self) -> str:
274 """
275 Find the longest common prefix of all words in the trie.
276 Returns empty string if no common prefix or trie is empty.
277 """
278 if self.is_empty():
279 return ""
280
281 prefix = []
282 node = self.root
283
284 while len(node.children) == 1 and not node.end_of_word:
285 char = next(iter(node.children.keys()))
286 prefix.append(char)
287 node = node.children[char]
288
289 return "".join(prefix)
290
291 def levels_traversal(self) -> dict[int, list[str]]:
292 """Return characters grouped by level (breadth-first). Level 0 is root."""
293 levels: dict[int, list[str]] = {}
294 queue = [(self.root, 0)]
295
296 for node, level in queue:
297 if level not in levels:
298 levels[level] = []
299 levels[level].append(node.char if node.char else "[root]")
300
301 for char in sorted(node.children.keys()):
302 queue.append((node.children[char], level + 1))
303
304 return levels
305
306 def min_depth(self) -> int:
307 """Return minimum depth (shortest word length). -1 if empty."""
308 if self.is_empty():
309 return -1
310
311 min_lvl = -1
312
313 def _dfs(node: TrieNode, depth: int):
314 nonlocal min_lvl
315 if node.end_of_word:
316 if min_lvl == -1 or depth < min_lvl:
317 min_lvl = depth
318 for child in node.children.values():
319 _dfs(child, depth + 1)
320
321 _dfs(self.root, 0)
322 return min_lvl
323
324 def max_depth(self) -> int:
325 """Return maximum depth (longest word length). -1 if empty."""
326 if self.is_empty():
327 return -1
328
329 max_lvl = 0
330
331 def _dfs(node: TrieNode, depth: int):
332 nonlocal max_lvl
333 if node.end_of_word and depth > max_lvl:
334 max_lvl = depth
335 for child in node.children.values():
336 _dfs(child, depth + 1)
337
338 _dfs(self.root, 0)
339 return max_lvl
340
341
342# Example usage:
343if __name__ == "__main__":
344 trie = Trie()
345
346 # Insert words
347 words = ["apple", "app", "application", "apply", "banana", "band", "bandana"]
348 for word in words:
349 trie.insert(word)
350
351 # Search
352 print("Search 'apple':", trie.search("apple")) # True
353 print("Search 'app':", trie.search("app")) # True
354 print("Search 'appl':", trie.search("appl")) # False (not a complete word)
355 print("'banana' in trie:", "banana" in trie) # True
356
357 # Prefix matching
358 print("Starts with 'app':", trie.starts_with("app")) # True
359 print("Starts with 'xyz':", trie.starts_with("xyz")) # False
360 print("Words with prefix 'app':", trie.get_words_with_prefix("app"))
361 # ['app', 'apple', 'application', 'apply']
362
363 # Get all words
364 print("All words:", trie.get_all_words())
365 # ['app', 'apple', 'application', 'apply', 'banana', 'band', 'bandana']
366
367 # Statistics
368 print("Word count:", trie.count_words()) # 7
369 print("Node count:", trie.count_nodes())
370
371 # Traversals
372 print("Preorder:", trie.traverse(Trie.TraverseOrder.preorder))
373 print("Levelorder:", trie.traverse(Trie.TraverseOrder.levelorder))
374 print("By levels:", trie.levels_traversal())
375
376 # Delete
377 trie.delete("app")
378 print("After deleting 'app':", trie.get_all_words())
379 print("Search 'app' after delete:", trie.search("app")) # False
380 print("Search 'apple' after delete:", trie.search("apple")) # True (still exists)
381
382 # Longest common prefix
383 trie2 = Trie()
384 for word in ["flower", "flow", "flight"]:
385 trie2.insert(word)
386 print("Longest common prefix:", trie2.longest_common_prefix()) # "fl"
387