[파이썬 | Python] 트라이 (Trie) 자료구조

문자열은 항상 어렵다. KMP도 그렇고, digit으로 정렬하는 것도 그렇고, 알면 알수록 머리아파지는 분야.

그만큼 어렵게 만들면 훨씬 어렵게도 만들 수 있다는 이야기겠지.

 

오늘은 트라이를 공부했다. Radix tree / Prefix tree 라고도 불리는데, 한 단어의 접두사(접두어)를 모두 저장하고 있다. (해당 단어에 도달하기까지의 문자들을 저장한다)

donghoon 이라는 단어를 보면, dong 도 접두사가 될 수 있고, do 도 접두사가 될 수 있다. 트라이에서는 이 단어들이 서로 포함관계에 있다는 것을 알려준다.

 

트라이에 "app", "ant", "apple"이라는 단어들을 저장한다고 하자. 트라이에는 지금까지의 모든 단어의 자취를 저장한다고 했다.

단어의 각 글자마다, 존재하지 않으면 새로 만들고, 존재한다면 그 글자로 들어가는 동작을 수행한다. 

단어의 끝에는 표시를 해 둔다. 트라이의 루트로부터 글자를 따라가다가, 표시가 된 곳을 방문한다면, 지금까지의 자취가 곧 하나의 단어가 된다.

트라이에 app, ant, apple을 순서대로 넣었을 때의 구조. 파란 글씨는 해당 글자로 끝나는 단어가 있다는 의미.

 

사실 C++에서는 포인터를 이용해서 트라이의 구조를 만드는데, 나는 각 노드를 하나의 딕셔너리로 보고 트라이를 구성했다.

트라이의 루트를 하나 잡아주면서 초기화한다. 최초 노드는 비어있는 딕셔너리이다.

class Trie:
    def __init__(self):
        self.root = {}

 

트라이는 문자열의 길이 $N$에 대해서 삽입과 검색을 $O(N)$에 해낸다.

삽입은 문자열의 각 문자에 대해서 노드가 있으면 그 노드로 들어가고, 없으면 새로운 노드를 현재 노드의 자식으로 넣어준다.

def insert(self, s):
    cur_node = self.root # 처음 루트로 잡아준다
    for c in s: # 받은 문자열 하나하나에 대해서
        if c not in cur_node: # 현재 노드에 그 문자가 있으면 들어가고, 그렇지 않으면 새로 만든다
            cur_node[c] = {}
        cur_node = cur_node[c]
    cur_node["*"] = s # "*" 노드를 만들어 '단어의 끝' 표시를 해 준다.

def search(self, s):
    cur_node = self.root
    for c in s:
        if c in cur_node:
            cur_node = cur_node[c]
        else:
            return False
    return "*" in cur_node # "*" 가 해당 노드에 있으면 그 글자로 끝나는 단어가 있다는 것이다.

생각보다는 어렵지 않은데, 이게 자주 쓰이지 않는 이유는 공간복잡도가 꽤 나가기 때문이다. 계속해서 노드를 만들어나가기 때문에 많은 문자열을 담게 되면 문자의 개수만큼의 공간을 차지하게 될 것이다. (최악의 경우)

 

---

연습문제

 

- BOJ 5052: 전화번호 목록

사실 이 문제는 str 의 startswith 를 쓰면 더 빨리 풀리지만, 문제가 워낙 트라이스러워서 한 번 풀어보자.

한 번호가 다른 번호의 접두사가 되는 일이 있어서는 안 된다. 트라이로 insert 해주면서, 노드에 단어의 끝 표시가 있는지 체크한다.

이 때 insert 는 단어가 짧은 순서대로 해 줘야 한다! 단어가 긴 단어 먼저 트라이에 저장하면, 접두사로 해당되는 짧은 문자가 들어왔을 때 NO 를 출력하지 못한다.

 

- BOJ 14725: 개미굴

문자 하나가 트라이 노드를 구성한다는 고정관념을 깨야 한다. "APPLE"이나 "BANANA"처럼 하나의 단어가 트라이의 노드를 구성한다는 점을 알고 문제를 해결하면 쉽게 해결할 수 있다.

출력같은 경우에는 딕셔너리를 리스트로 변환해 노드를 알파벳 순으로 정렬한 뒤, 재귀적으로 출력하면 된다.

 

- BOJ 13502: 단어퍼즐 2

처음으로 백준에 200000bytes 가 넘는 코드를 제출하게 한 문제다. 컴퓨터 성능도 좋지 않아서 붙여넣기로 코드 제출하다가 터진 게 두어 번. 

해답은 주어진 사전 파일을 모두 트라이에 저장한 다음, 백트래킹을 사용해서 모든 칸에 대해서 단어를 찾으면 된다. 중복된 단어를 찾을 수 있기 때문에, set 을 사용해서 단어를 넣어주고, 집합의 크기를 출력했다.

 

- BOJ 16934: 게임 닉네임

새로운 노드를 만든다면, 지금까지 있던 닉네임이 아니므로 처음 새로운 노드를 만드는 시점까지만 출력하면 된다.

insert 메서드에서 새로운 노드를 만들게 된 시점을 기록하고, 그 인덱스까지 닉네임을 출력한다. 만일 같은 닉네임이 있다면, 숫자를 더해서 카운트해주는 것도 잊지 말자.