[문자열] 트라이(Trie) (Python)

트라이란?

트라이(Trie)란 주로 문자열을 효율적으로 저장, 탐색하기 위한 트리 형태의 자료 구조를 의미한다.

트라이의 예시 (출처=위키백과)

트라이는 문자열의 맨 앞 문자부터 계층적으로 저장하고 있다.
위 그림에서 'tea'를 찾으려면 순서대로 't', 'te', 'tea'로 찾아 내려오면 된다.
이처럼 트라이에서 찾으려는 문자열의 길이가 $m$일 때, $O(m)$의 시간 만에 원하는 문자열을 찾을 수 있다.

 

구현해보기

삽입

'tax'라는 문자열을 삽입해보자.

1. root의 자식 중 't'가 있으므로 해당 노드로 이동한다.
2. 't'의 자식 중 'a'('ta')가 없으므로 해당 노드를 생성 후 이동한다.
3. 'a'의 자식 중 'x'('tax')가 없으므로 해당 노드를 생성 후 이동한다.

해당 구조를 코드로 작성하면 아래와 같다.

class Node:
    def __init__(self, key, data = None):
        self.key = key # 해당 노드에 대응되는 문자 저장
        self.data = data # 해당 노드로 최종 완성되는 문자열 저장
        self.children = {}


class Trie:
    def __init__(self):
        self.root = Node(None)

    def insert(self, string):
        now = self.root

        for char in string:
            if char not in now.children:
                now.children[char] = Node(char)
            now = now.children[char]

        now.data = string

 

탐색

탐색도 삽입과 크게 다르지 않다.
탐색을 원하는 문자열의 맨 앞 문자를 root부터 하위 노드로 쭉 비교해나가면 된다.

    def search(self, string):
        now = self.root

        for char in string:
            if char in now.children:
                now = now.children[char]
            else:
                return False

        if now.data:
            return True

 

관련 문제 (BOJ)

• [14725] 개미굴: 트라이 자료 구조를 맛볼 수 있는 문제
• [16934] 게임 닉네임: 트라이 활용 문제
• [5670] 휴대폰 자판: 트라이의 코드를 변형해서 해결해야 하는 문제
• [16903] 수열과 쿼리 20: 문자열이 아닌 비트로 트라이를 만들어 해결하는 문제