LeetCode 208 - Implement Trie (Prefix Tree)

이 문제는 문자열 집합에 대해 단어 삽입, 완전한 단어 검색, prefix 검색을 구현하는 Medium 자료구조 문제다. 입력 문자가 소문자 영어 알파벳으로 제한되어 있어서, 각 노드가 26칸짜리 자식 배열을 갖고 단어의 끝은 is_end 플래그로 표시하는 Trie로 풀었다.

문제 링크 & 설명

문제 링크: 208. Implement Trie (Prefix Tree)
요약: Trie 클래스를 만들고, insert, search, startsWith 세 가지 연산을 구현하는 문제다.

Trie는 문자열을 문자 단위로 저장하는 tree 자료구조다. 일반적인 set이나 hash table이 단어 전체를 하나의 key로 보는 데 비해, Trie는 같은 prefix를 공유하는 단어들이 같은 경로를 나눠 쓴다.

예를 들어 apple을 넣으면 root에서 시작해 a, p, p, l, e로 내려가는 경로가 만들어진다. 이 상태에서 app을 검색하면 경로 자체는 존재하지만, app이라는 단어를 넣은 적은 없다. 그래서 search는 false여야 하고, startsWith는 true여야 한다.

insert("apple")
search("apple")    -> true
search("app")      -> false
startsWith("app")  -> true
insert("app")
search("app")      -> true

이 차이 때문에 노드가 존재하는지만으로는 단어 검색을 끝낼 수 없다. 어떤 노드가 실제로 삽입된 단어의 마지막 문자였는지를 따로 저장해야 한다.

접근 방법

먼저 TrieNode를 따로 만들고, 각 노드가 children 배열과 is_end 플래그를 갖도록 했다.

class TrieNode:
    def __init__(self):
        self.children = [None] * 26
        self.is_end = False

children 배열은 소문자 알파벳 26개에 대응한다. 문자 c가 들어오면 ord(c) - ord(‘a’)로 인덱스를 계산하고, 해당 위치에 자식 노드가 없으면 새 TrieNode를 만든다. 알파벳 크기가 고정되어 있으니 dictionary를 쓰지 않고 배열로 바로 접근할 수 있다.

insert는 root에서 시작해 단어의 각 문자를 따라 내려간다. 경로가 없으면 만들고, 있으면 기존 노드를 그대로 재사용한다. 마지막 문자까지 처리한 뒤 현재 노드의 is_end를 True로 바꾸면, 그 경로가 하나의 완전한 단어로 등록된다.

cur = self.root

for c in word:
    idx = ord(c) - ord('a')
    if not cur.children[idx]:
        cur.children[idx] = TrieNode()
    cur = cur.children[idx]

cur.is_end = True

search도 같은 방식으로 문자를 하나씩 따라 내려간다. 중간에 필요한 자식 노드가 없으면 그 단어는 Trie에 없으므로 바로 False를 반환한다. 모든 문자를 끝까지 따라간 뒤에는 현재 노드의 is_end를 확인한다.

여기서 search와 startsWith가 갈린다. search는 입력 문자열 전체가 삽입된 단어인지 확인해야 하므로 마지막에 is_end가 True여야 한다. 반대로 startsWith는 입력 prefix에 해당하는 경로가 존재하는지만 확인하면 되기 때문에, 마지막 노드의 is_end를 보지 않고 True를 반환한다.

apple만 삽입한 상태에서 app을 생각하면 이 차이가 바로 드러난다. a, p, p 경로는 있으므로 startsWith(“app”)은 True지만, 세 번째 p 노드의 is_end는 아직 False라서 search(“app”)은 False다. 이후 app을 insert하면 같은 경로를 다시 내려간 뒤 그 노드의 is_end만 True로 바뀐다.

복잡도 분석

단어 길이를 L, prefix 길이를 P라고 하자.

insert 시간 복잡도: O(L)
- 단어의 각 문자를 한 번씩 처리한다. 각 문자마다 배열 인덱스 계산과 자식 노드 접근은 고정 알파벳 크기 안에서 상수 시간이다.
insert 추가 공간 복잡도: O(L)
- 이미 존재하는 경로는 재사용하지만, 공유되는 prefix가 전혀 없다면 문자마다 새 노드를 만들 수 있다. 새 노드 하나는 26칸 children 배열과 is_end 플래그를 가진다. 알파벳 크기 26은 고정값이므로, 새로 할당되는 노드 수 기준으로는 O(L)이다.
search 시간 복잡도: O(L)
- 검색할 단어의 문자를 순서대로 따라가므로 입력 길이만큼 탐색한다.
search 추가 공간 복잡도: O(1)
- 현재 노드를 가리키는 포인터와 인덱스만 사용한다.
startsWith 시간 복잡도: O(P)
- prefix의 각 문자를 한 번씩 따라간다.
startsWith 추가 공간 복잡도: O(1)
- search와 마찬가지로 별도 자료구조를 만들지 않는다.

여러 단어를 계속 삽입한 뒤 Trie 전체가 차지하는 저장 공간은 삽입된 문자열들이 만든 노드 수에 비례한다. 전체 삽입 문자 수를 T라고 하면 최악의 경우 노드 수는 O(T)이고, 각 노드는 26개의 자식 참조를 고정으로 가진다. 그래서 실제 상수 비용은 꽤 큰 편이지만, 알파벳 크기가 문제 조건에서 고정되어 있으므로 보통 O(T) 공간으로 정리한다.

구현 코드

class TrieNode:
    def __init__(self):
        self.children = [None] * 26
        self.is_end = False


class Trie:

    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        cur = self.root

        for c in word:
            idx = ord(c) - ord('a')
            if not cur.children[idx]:
                cur.children[idx] = TrieNode()
            cur = cur.children[idx]

        cur.is_end = True

    def search(self, word: str) -> bool:
        cur = self.root

        for c in word:
            idx = ord(c) - ord('a')
            if cur.children[idx]:
                cur = cur.children[idx]
            else:
                return False

        return cur.is_end

    def startsWith(self, prefix: str) -> bool:
        cur = self.root

        for c in prefix:
            idx = ord(c) - ord('a')
            if cur.children[idx]:
                cur = cur.children[idx]
            else:
                return False

        return True

요약 및 회고

이 문제에서 핵심은 TrieNode의 모양을 먼저 정하는 일이었다. 입력 문자가 소문자 영어로 고정되어 있으니 26칸 배열을 둘 수 있고, 같은 prefix를 공유하는 단어들은 같은 노드 경로를 재사용한다. 여기에 is_end 플래그를 붙이면 경로 존재 여부와 단어 종료 여부를 분리할 수 있다.

search와 startsWith는 거의 같은 탐색 코드를 쓰지만, 마지막 판단 기준이 다르다. search는 완전한 단어를 찾는 연산이라 is_end까지 확인해야 하고, startsWith는 prefix 경로만 확인하면 된다. 겉으로는 작은 차이지만, apple을 넣은 뒤 app을 찾는 예시처럼 Trie의 의미를 결정하는 부분은 바로 그 마지막 한 줄이었다.