Python으로 트라이(Trie, 접두사 트리) 자료구조 구현하기

트라이(Trie)는 문자열을 효율적으로 저장하고 탐색하기 위한 트리 기반 자료구조로, '접두사 트리'라고도 불립니다. 이번 글에서는 Python으로 insert(), search(), startsWith() 세 가지 기본 연산을 갖춘 트라이를 구현하는 방법을 살펴보겠습니다. 편의상 모든 입력은 소문자 영어로만 이루어져 있다고 가정합니다.

동작 예시

다음과 같이 메서드를 호출했을 때 각 연산의 결과를 확인할 수 있습니다.

trie = Trie()
trie.insert("apple")
trie.search("apple") → True 반환 (단어 "apple"이 존재)
trie.search("app") → False 반환 ("app"은 저장된 단어가 아님)
trie.startsWith("app") → True 반환 ("apple"이 "app"으로 시작)
trie.insert("app")
trie.search("app") → True 반환 (이제 "app"도 저장된 단어)

구현 전략

이 문제는 중첩된 딕셔너리(nested dictionary)를 활용하면 간단하게 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

1. 초기화

child라는 이름의 딕셔너리를 하나 생성합니다. 이 딕셔너리가 트라이의 루트 노드 역할을 하며, 각 자식 노드 역시 문자를 키로 갖는 딕셔너리 형태로 표현됩니다.

2. insert(word) 메서드

current를 루트(child)로 설정합니다.
단어의 각 문자 l에 대해:
- l이 current에 없으면 current[l] = {}로 새 딕셔너리를 만듭니다.
- current = current[l]로 한 단계 내려갑니다.
모든 문자를 처리한 후 마지막 노드에 current['#'] = 1을 설정하여 단어의 끝임을 표시합니다.

3. search(word) 메서드

current를 루트로 설정한 뒤 단어의 각 문자를 따라 내려갑니다.
중간에 해당 문자가 존재하지 않으면 즉시 False를 반환합니다.
탐색을 마친 후 마지막 노드에 '#' 키가 있으면(완성된 단어라면) True, 없으면 False를 반환합니다.

4. startsWith(prefix) 메서드

search와 동일하게 접두사의 각 문자를 따라 내려갑니다.
경로가 끊기면 False를 반환하고, 끝까지 도달하면 그 노드가 완성된 단어인지와 관계없이 True를 반환합니다.

search와 startsWith의 차이는 바로 이 부분입니다. search는 정확히 저장된 단어인지까지 확인하는 반면, startsWith는 해당 접두사로 시작하는 단어의 존재 가능 여부만 판단합니다.

전체 코드

실제 구현 코드를 통해 더 잘 이해해 보겠습니다.

class Trie(object):
    def __init__(self):
        self.child = {}
    def insert(self, word):
        current = self.child
        for l in word:
            if l not in current:
                current[l] = {}
            current = current[l]
        current['#'] = 1
    def search(self, word):
        current = self.child
        for l in word:
            if l not in current:
                return False
            current = current[l]
        return '#' in current
    def startsWith(self, prefix):
        current = self.child
        for l in prefix:
            if l not in current:
                return False
            current = current[l]
        return True

ob1 = Trie()
ob1.insert("apple")
print(ob1.search("apple"))
print(ob1.search("app"))
print(ob1.startsWith("app"))
ob1.insert("app")
print(ob1.search("app"))

실행 결과

True
False
True
True

마무리

이처럼 딕셔너리만 활용해도 별도의 노드 클래스 없이 트라이를 깔끔하게 구현할 수 있습니다. 삽입, 검색, 접두사 확인 모두 단어 길이에 비례하는 O(L) 시간 복잡도를 가지므로, 자동 완성 기능이나 대량의 문자열 집합을 다루는 문제에서 매우 유용하게 사용됩니다.