트라이(Trie)는 문자열을 효율적으로 탐색하기 위한 자료구조로, 문자열을 선형 시간 안에 찾아낼 수 있다. 트라이에 word를 삽입하는 insert 메소드와 prefix로 시작하는 문자열들의 리스트를 반환하는 starts_with 메소드를 구현했다.
class Node:
def __init__(self, char):
self.char = char
self.word = None
self.children = {}
self.num_words = 0
class Trie():
def __init__(self):
self.root = Node(None)
def insert(self, word):
cur_node = self.root
for char in word:
if char not in cur_node.children:
cur_node.children[char] = Node(char)
cur_node = cur_node.children[char]
cur_node.num_words += 1
cur_node.word = word
def starts_with(self, prefix):
from collections import deque
cur_node = self.root
for char in prefix:
if char not in cur_node.children:
return []
cur_node = cur_node.children[char]
# BFS
queue = deque([cur_node]); result = [];
while queue:
node = queue.popleft()
if node.word:
result.append(node.word)
for next_node in node.children.values():
queue.append(next_node)
return result