1"""
2νΈλΌμ΄ (Trie / Prefix Tree)
3Trie Data Structure
4
5λ¬Έμμ΄ κ²μκ³Ό μ λμ¬ κΈ°λ° μ°μ°μ μν νΈλΌμ΄λ₯Ό ꡬνν©λλ€.
6"""
7
8from typing import List, Optional, Dict
9from collections import defaultdict
10
11
12# =============================================================================
13# 1. κΈ°λ³Έ νΈλΌμ΄ (λμ
λ리 κΈ°λ°)
14# =============================================================================
15
16class TrieNode:
17 """νΈλΌμ΄ λ
Έλ"""
18
19 def __init__(self):
20 self.children: Dict[str, 'TrieNode'] = {}
21 self.is_end: bool = False
22 self.count: int = 0 # μ΄ μ λμ¬λ‘ μμνλ λ¨μ΄ μ
23
24
25class Trie:
26 """
27 νΈλΌμ΄ (μ λμ¬ νΈλ¦¬)
28 - μ½μ
: O(m), m = λ¨μ΄ κΈΈμ΄
29 - κ²μ: O(m)
30 - μ λμ¬ κ²μ: O(m)
31 """
32
33 def __init__(self):
34 self.root = TrieNode()
35
36 def insert(self, word: str) -> None:
37 """λ¨μ΄ μ½μ
- O(m)"""
38 node = self.root
39
40 for char in word:
41 if char not in node.children:
42 node.children[char] = TrieNode()
43 node = node.children[char]
44 node.count += 1
45
46 node.is_end = True
47
48 def search(self, word: str) -> bool:
49 """λ¨μ΄ μ‘΄μ¬ μ¬λΆ κ²μ - O(m)"""
50 node = self._find_node(word)
51 return node is not None and node.is_end
52
53 def starts_with(self, prefix: str) -> bool:
54 """μ λμ¬λ‘ μμνλ λ¨μ΄ μ‘΄μ¬ μ¬λΆ - O(m)"""
55 return self._find_node(prefix) is not None
56
57 def count_prefix(self, prefix: str) -> int:
58 """μ λμ¬λ‘ μμνλ λ¨μ΄ κ°μ - O(m)"""
59 node = self._find_node(prefix)
60 return node.count if node else 0
61
62 def _find_node(self, prefix: str) -> Optional[TrieNode]:
63 """μ λμ¬μ ν΄λΉνλ λ
Έλ μ°ΎκΈ°"""
64 node = self.root
65
66 for char in prefix:
67 if char not in node.children:
68 return None
69 node = node.children[char]
70
71 return node
72
73 def delete(self, word: str) -> bool:
74 """λ¨μ΄ μμ - O(m)"""
75
76 def _delete(node: TrieNode, word: str, depth: int) -> bool:
77 if depth == len(word):
78 if not node.is_end:
79 return False
80 node.is_end = False
81 return len(node.children) == 0
82
83 char = word[depth]
84 if char not in node.children:
85 return False
86
87 should_delete = _delete(node.children[char], word, depth + 1)
88
89 if should_delete:
90 del node.children[char]
91 return len(node.children) == 0 and not node.is_end
92
93 node.children[char].count -= 1
94 return False
95
96 return _delete(self.root, word, 0)
97
98
99# =============================================================================
100# 2. μλμμ± (Autocomplete)
101# =============================================================================
102
103class AutocompleteSystem:
104 """μλμμ± μμ€ν
"""
105
106 def __init__(self):
107 self.root = TrieNode()
108
109 def add_word(self, word: str, weight: int = 1) -> None:
110 """λ¨μ΄ μΆκ° (κ°μ€μΉ ν¬ν¨)"""
111 node = self.root
112
113 for char in word:
114 if char not in node.children:
115 node.children[char] = TrieNode()
116 node = node.children[char]
117
118 node.is_end = True
119 node.count = weight # μ¬κΈ°μ countλ λΉλ/κ°μ€μΉ
120
121 def autocomplete(self, prefix: str, limit: int = 5) -> List[str]:
122 """μ λμ¬λ‘ μμνλ λ¨μ΄ λͺ©λ‘ - O(m + k), k = κ²°κ³Ό μ"""
123 node = self.root
124
125 # μ λμ¬ λ
Έλ μ°ΎκΈ°
126 for char in prefix:
127 if char not in node.children:
128 return []
129 node = node.children[char]
130
131 # DFSλ‘ λͺ¨λ λ¨μ΄ μμ§
132 results = []
133 self._collect_words(node, prefix, results)
134
135 # λΉλμ μ λ ¬ ν μμ limitκ°
136 results.sort(key=lambda x: x[1], reverse=True)
137 return [word for word, _ in results[:limit]]
138
139 def _collect_words(self, node: TrieNode, prefix: str, results: List) -> None:
140 """λ
Έλμμ λͺ¨λ μμ± λ¨μ΄ μμ§"""
141 if node.is_end:
142 results.append((prefix, node.count))
143
144 for char, child in node.children.items():
145 self._collect_words(child, prefix + char, results)
146
147
148# =============================================================================
149# 3. μμΌλμΉ΄λ κ²μ
150# =============================================================================
151
152class WildcardTrie:
153 """μμΌλμΉ΄λ '.'λ₯Ό μ§μνλ νΈλΌμ΄"""
154
155 def __init__(self):
156 self.root = TrieNode()
157
158 def add_word(self, word: str) -> None:
159 node = self.root
160 for char in word:
161 if char not in node.children:
162 node.children[char] = TrieNode()
163 node = node.children[char]
164 node.is_end = True
165
166 def search(self, word: str) -> bool:
167 """'.'λ μμμ ν λ¬Έμμ 맀μΉ"""
168 return self._search(self.root, word, 0)
169
170 def _search(self, node: TrieNode, word: str, index: int) -> bool:
171 if index == len(word):
172 return node.is_end
173
174 char = word[index]
175
176 if char == '.':
177 # λͺ¨λ μμ λ
Έλ μλ
178 for child in node.children.values():
179 if self._search(child, word, index + 1):
180 return True
181 return False
182 else:
183 if char not in node.children:
184 return False
185 return self._search(node.children[char], word, index + 1)
186
187
188# =============================================================================
189# 4. μ΅μ₯ κ³΅ν΅ μ λμ¬ (LCP)
190# =============================================================================
191
192def longest_common_prefix(words: List[str]) -> str:
193 """
194 λ¨μ΄ λ°°μ΄μ μ΅μ₯ κ³΅ν΅ μ λμ¬
195 νΈλΌμ΄ νμ©
196 """
197 if not words:
198 return ""
199
200 # νΈλΌμ΄ ꡬμ±
201 trie = Trie()
202 for word in words:
203 trie.insert(word)
204
205 # 루νΈμμ λΆκΈ°μ κΉμ§ μ΄λ
206 prefix = []
207 node = trie.root
208
209 while node:
210 # μμμ΄ νλμ΄κ³ μ’
λ£ λ
Έλκ° μλ λλ§ κ³μ
211 if len(node.children) == 1 and not node.is_end:
212 char = list(node.children.keys())[0]
213 prefix.append(char)
214 node = node.children[char]
215 else:
216 break
217
218 return ''.join(prefix)
219
220
221# =============================================================================
222# 5. λ¨μ΄ μ¬μ (Word Dictionary)
223# =============================================================================
224
225class WordDictionary:
226 """
227 λ¨μ΄ μΆκ°/κ²μ μ¬μ
228 - μ νν κ²μ
229 - μ λμ¬ κ²μ
230 - μμΌλμΉ΄λ κ²μ
231 """
232
233 def __init__(self):
234 self.trie = Trie()
235 self.wildcard_trie = WildcardTrie()
236
237 def add_word(self, word: str) -> None:
238 self.trie.insert(word)
239 self.wildcard_trie.add_word(word)
240
241 def search_exact(self, word: str) -> bool:
242 """μ νν λ¨μ΄ κ²μ"""
243 return self.trie.search(word)
244
245 def search_prefix(self, prefix: str) -> bool:
246 """μ λμ¬ κ²μ"""
247 return self.trie.starts_with(prefix)
248
249 def search_pattern(self, pattern: str) -> bool:
250 """ν¨ν΄ κ²μ (. = μμ λ¬Έμ)"""
251 return self.wildcard_trie.search(pattern)
252
253
254# =============================================================================
255# 6. XOR νΈλΌμ΄ (μ΅λ XOR μ μ°ΎκΈ°)
256# =============================================================================
257
258class XORTrie:
259 """
260 λΉνΈ νΈλΌμ΄λ‘ μ΅λ XOR μ μ°ΎκΈ°
261 κ° μ«μλ₯Ό μ΄μ§μλ‘ μ μ₯
262 """
263
264 def __init__(self, max_bits: int = 31):
265 self.root = {}
266 self.max_bits = max_bits
267
268 def insert(self, num: int) -> None:
269 """μ«μ μ½μ
- O(max_bits)"""
270 node = self.root
271
272 for i in range(self.max_bits, -1, -1):
273 bit = (num >> i) & 1
274 if bit not in node:
275 node[bit] = {}
276 node = node[bit]
277
278 def find_max_xor(self, num: int) -> int:
279 """numκ³Ό XORνμ λ μ΅λκ° λλ κ° λ°ν - O(max_bits)"""
280 node = self.root
281 result = 0
282
283 for i in range(self.max_bits, -1, -1):
284 bit = (num >> i) & 1
285 # XOR μ΅λνλ₯Ό μν΄ λ°λ λΉνΈ μ ν
286 toggle_bit = 1 - bit
287
288 if toggle_bit in node:
289 result |= (1 << i)
290 node = node[toggle_bit]
291 elif bit in node:
292 node = node[bit]
293 else:
294 break
295
296 return result
297
298
299def find_maximum_xor(nums: List[int]) -> int:
300 """
301 λ°°μ΄μμ λ μμ XOR μ΅λκ°
302 μκ°λ³΅μ‘λ: O(n * max_bits)
303 """
304 if len(nums) < 2:
305 return 0
306
307 xor_trie = XORTrie()
308 max_xor = 0
309
310 for num in nums:
311 xor_trie.insert(num)
312 max_xor = max(max_xor, xor_trie.find_max_xor(num))
313
314 return max_xor
315
316
317# =============================================================================
318# 7. μ€μ λ¬Έμ : λ¨μ΄ κ²μ II (Word Search II)
319# =============================================================================
320
321def find_words(board: List[List[str]], words: List[str]) -> List[str]:
322 """
323 2D 보λμμ λ¨μ΄ λͺ©λ‘ μ°ΎκΈ°
324 νΈλΌμ΄ + DFS
325 """
326 # νΈλΌμ΄ ꡬμ±
327 trie = Trie()
328 for word in words:
329 trie.insert(word)
330
331 rows, cols = len(board), len(board[0])
332 result = set()
333
334 def dfs(r: int, c: int, node: TrieNode, path: str) -> None:
335 if node.is_end:
336 result.add(path)
337
338 if r < 0 or r >= rows or c < 0 or c >= cols:
339 return
340
341 char = board[r][c]
342 if char == '#' or char not in node.children:
343 return
344
345 board[r][c] = '#' # λ°©λ¬Έ νμ
346 next_node = node.children[char]
347
348 for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
349 dfs(r + dr, c + dc, next_node, path + char)
350
351 board[r][c] = char # 볡μ
352
353 for r in range(rows):
354 for c in range(cols):
355 dfs(r, c, trie.root, "")
356
357 return list(result)
358
359
360# =============================================================================
361# 8. μ λμ¬/μ λ―Έμ¬ λμ κ²μ
362# =============================================================================
363
364class PrefixSuffixTrie:
365 """μ λμ¬μ μ λ―Έμ¬λ₯Ό λμμ κ²μνλ νΈλΌμ΄"""
366
367 def __init__(self, words: List[str]):
368 self.trie = {}
369
370 for idx, word in enumerate(words):
371 # λͺ¨λ μ λ―Έμ¬#λ¨μ΄ ννλ‘ μ μ₯
372 key = word + '#' + word
373 for i in range(len(word) + 1):
374 node = self.trie
375 for char in key[i:]:
376 if char not in node:
377 node[char] = {'idx': -1}
378 node = node[char]
379 node['idx'] = idx # κ°μ₯ ν° μΈλ±μ€ μ μ₯
380
381 def search(self, prefix: str, suffix: str) -> int:
382 """μ λμ¬μ μ λ―Έμ¬ λͺ¨λ λ§μ‘±νλ λ¨μ΄μ μΈλ±μ€"""
383 key = suffix + '#' + prefix
384 node = self.trie
385
386 for char in key:
387 if char not in node:
388 return -1
389 node = node[char]
390
391 return node.get('idx', -1)
392
393
394# =============================================================================
395# ν
μ€νΈ
396# =============================================================================
397
398def main():
399 print("=" * 60)
400 print("νΈλΌμ΄ (Trie / Prefix Tree) μμ ")
401 print("=" * 60)
402
403 # 1. κΈ°λ³Έ νΈλΌμ΄
404 print("\n[1] κΈ°λ³Έ νΈλΌμ΄")
405 trie = Trie()
406 words = ["apple", "app", "application", "banana", "band"]
407 for word in words:
408 trie.insert(word)
409 print(f" μ½μ
: {words}")
410 print(f" search('app'): {trie.search('app')}")
411 print(f" search('application'): {trie.search('application')}")
412 print(f" search('apply'): {trie.search('apply')}")
413 print(f" starts_with('app'): {trie.starts_with('app')}")
414 print(f" count_prefix('app'): {trie.count_prefix('app')}")
415
416 # 2. μλμμ±
417 print("\n[2] μλμμ±")
418 auto = AutocompleteSystem()
419 search_words = [("hello", 5), ("help", 3), ("helicopter", 2), ("hero", 4), ("world", 1)]
420 for word, weight in search_words:
421 auto.add_word(word, weight)
422 print(f" λ¨μ΄/λΉλ: {search_words}")
423 print(f" 'hel' μλμμ±: {auto.autocomplete('hel')}")
424 print(f" 'he' μλμμ±: {auto.autocomplete('he')}")
425
426 # 3. μμΌλμΉ΄λ κ²μ
427 print("\n[3] μμΌλμΉ΄λ κ²μ")
428 wild = WildcardTrie()
429 for word in ["bad", "dad", "mad", "pad"]:
430 wild.add_word(word)
431 print(f" λ¨μ΄: ['bad', 'dad', 'mad', 'pad']")
432 print(f" search('.ad'): {wild.search('.ad')}")
433 print(f" search('b..'): {wild.search('b..')}")
434 print(f" search('..d'): {wild.search('..d')}")
435 print(f" search('b.d'): {wild.search('b.d')}")
436
437 # 4. μ΅μ₯ κ³΅ν΅ μ λμ¬
438 print("\n[4] μ΅μ₯ κ³΅ν΅ μ λμ¬")
439 words_lcp = ["flower", "flow", "flight"]
440 lcp = longest_common_prefix(words_lcp)
441 print(f" λ¨μ΄: {words_lcp}")
442 print(f" LCP: '{lcp}'")
443
444 # 5. XOR νΈλΌμ΄
445 print("\n[5] μ΅λ XOR (XOR νΈλΌμ΄)")
446 nums = [3, 10, 5, 25, 2, 8]
447 max_xor = find_maximum_xor(nums)
448 print(f" λ°°μ΄: {nums}")
449 print(f" μ΅λ XOR: {max_xor}")
450 print(f" (5 XOR 25 = {5 ^ 25})")
451
452 # 6. λ¨μ΄ μ¬μ
453 print("\n[6] λ¨μ΄ μ¬μ ")
454 dictionary = WordDictionary()
455 for word in ["hello", "help", "world"]:
456 dictionary.add_word(word)
457 print(f" λ¨μ΄: ['hello', 'help', 'world']")
458 print(f" exact 'help': {dictionary.search_exact('help')}")
459 print(f" prefix 'hel': {dictionary.search_prefix('hel')}")
460 print(f" pattern 'h.l.o': {dictionary.search_pattern('h.l.o')}")
461
462 # 7. λ¨μ΄ μμ
463 print("\n[7] λ¨μ΄ μμ ")
464 trie2 = Trie()
465 for word in ["apple", "app"]:
466 trie2.insert(word)
467 print(f" μ½μ
: ['apple', 'app']")
468 print(f" search('app'): {trie2.search('app')}")
469 trie2.delete("app")
470 print(f" delete('app') ν")
471 print(f" search('app'): {trie2.search('app')}")
472 print(f" search('apple'): {trie2.search('apple')}")
473
474 print("\n" + "=" * 60)
475
476
477if __name__ == "__main__":
478 main()