파이썬 Trie

Trie

백준 개미굴 문제를 풀다 답을 보는 과정에서 Trie에 대해 알게되어 공부할겸 기록하고자 한다.

Trie(Tree)의 개념과 특징

트라이(Trie)는 문자열을 효율적으로 저장하고 검색하는 트리 형태의 자료구조입니다. 주로 자동완성 기능이나 사전 검색과 같이 문자열을 빠르게 탐색하는데 사용됩니다. 이 자료구조는 라데릭스 트리(radix tree), 접두사 트리(prefix tree), 탐색 트리(retrieval tree) 등으로도 불립니다.

트라이는 동적인 집합이나 연관 배열을 저장하기 위한 트리 구조로, 문자열을 일일이 비교하지 않고도 노드들을 탐색함으로써 검색 속도를 크게 향상시킵니다. 이로 인해 문자열 탐색과 관련된 작업에 매우 효과적입니다. 하지만 각 노드가 자식들에 대한 포인터를 배열로 저장하기 때문에 저장 공간의 사용량이 큰 단점이 있습니다.

시간 복잡도

제일 긴 문자열의 길이를 m, 총 문자열들의 수를 n이라 할 때

• 생성시 시간 복잡도 : O(m*n)
• 삽입 시간 복잡도 : O(m)
• 탐색시 시간 복잡도 : O(m)

파이썬으로 Trie 구현하기

class Node:
    def __init__(self):
        self.children = {}  # 자식 노드를 저장하는 딕셔너리 (문자: 노드)
        self.is_end_of_word = False  # 해당 노드에서 문자열이 끝나는지 여부


class Trie:
    def __init__(self):
        self.root = Node()  # 루트 노드 생성

    def insert(self, word):
        """
        Trie에 문자열을 삽입하는 메서드
        """
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = Node()  # 없는 문자일 경우 노드 생성
            node = node.children[char]
        node.is_end_of_word = True  # 문자열의 끝을 표시

    def search(self, word):
        """
        Trie에서 문자열을 검색하는 메서드
        """
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]
        return node.is_end_of_word  # 문자열이 끝까지 탐색되면 해당 노드의 is_end_of_word 반환

    def starts_with(self, prefix):
        """
        Trie에서 주어진 접두사로 시작하는 문자열을 검색하는 메서드
        """
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        return True


# Trie 사용 예시
trie = Trie()
trie.insert("apple")
trie.insert("app")
print(trie.search("apple"))  # 출력: True
print(trie.search("app"))    # 출력: True
print(trie.search("ap"))     # 출력: False
print(trie.starts_with("app"))  # 출력: True