PS 문자열 알고리즘 - 8. 1D Trie 구현

이 글은 독자가 trie 자료 구조에 대해 이미 알고 있다는 전제 하에서 시작한다. trie라는 자료 구조를 모른다면 이 글을 읽기 전에 다른 블로그를 읽고 https://www.acmicpc.net/problem/14725와 같은 문제를 해결해본 후 이 글을 읽는 것을 권장한다.

---

 Trie라는 자료 구조를 일반적으로 구현할 때에는 Node 클래스를 만들고 여기에 딕셔너리 혹은 리스트가 들어가서 여기에 다시 Node가 저장되는 형태로 구현한다. 하지만 이렇게 리스트나 딕셔너리를 많이 만들 경우 메모리 초과가 나기 십상이며, 접근 자체에도 긴 시간이 걸린다.

 

 이를 개선하는 대표적인 방법으로, Trie 자료 구조를 1차원 배열을 이용해서 구현하는 방식이 있다. 우리가 일반적으로 알고 있는 트라이의 구조에 인덱스를 매겨서 그림으로 그려보면 아래 그림과 같다.

 

 

오른쪽에 적은 대로 HE, SHE, SO, SHY라는 문자열을 순서대로 저장한 트라이이다. 이때 파란색 숫자와 같이 루트 노드를 0으로 한 후 노드가 생길 때마다 1씩 증가하면서 번호를 붙여주고, 번호가 $k$인 노드의 자식 노드들에 접근하기 위한 인덱스 $c$는 1차원 리스트 내에서 $26k \le c \le 26k + 25$에 속하게 만드는 것이다.

 

이때 26은 알파벳 대문자의 개수, 즉 가능한 알파벳(character)의 개수를 의미한다. 또한 리스트의 인덱스가 초과되면 안되므로 1차원 리스트는 알파벳의 개수와 노드의 최대 개수를 곱한 만큼 필요하고, 노드의 최대 개수는 다시 추가할 문자열의 수와 문자열의 길이의 최댓값으로 구해야 한다. 또한, A ~ Z 각각을 0에서 25까지의 인덱스에 대응해야 하므로 이를 대응하기 위한 함수도 필요하다.

 

또한, 각 노드를 지나간 문자열의 수를 저장하는 cnt 배열이나 해당 노드를 끝으로 하는 문자열이 있는지 저장하는 end 배열 등을 추가로 선언하여 Trie의 init 함수를 구현할 수 있다. 이를 구현한 것이 아래와 같다.

def __init__(self, alpha, leng, maxword, func):
    # alpha: the number of characters
    # max_node: the maximum number of nodes
    # func: function that translates characters to index num
    self.alpha = alpha
    self.maxnode = leng * maxword
    self.trie = [0] * (alpha * self.max_node)
    self.cnt = [0] * self.max_node
    self.end = [0] * self.max_node
    self.num = 0
    self.f = func

 

이번에는 가장 중요한 insert 함수를 구현할 것이다. 먼저, 구현한 코드를 살펴본 후 이에 대한 해설을 덧붙이겠다.

def insert(self, word):
    i = 0
    self.cnt[i] += 1
    for ch in word:
        idx = self.f(ch)
        pt = self.alpha * i + idx
        if self.trie[pt] == 0:
            self.num += 1
            self.trie[pt] = self.num
        i = self.trie[pt]
        self.cnt[i] += 1
    self.end[i] += 1

 

$i$라는 변수가 가장 핵심적인 역할을 하는 변수로서 현재 지나고 있는 노드의 번호를 저장하고 있다. 처음에는 루트 노드에서 출발하므로 $i = 0$이며, word의 앞 글자부터 살피며 ch를 idx로 trie 생성 시 만든 함수를 이용해서 변환하고, pt가 다음 자식 노드의 번호를 가리킨다. 만약, 이미 지난 자식 노드가 아닌 경우 self.trie[pt]의 값이 0이 되고, 이때는 self.num을 1 증가시켜서 trie라는 1차원 배열에 26칸의 공간을 추가로 이 노드에 할당해준다. 그렇게 해서 i를 갱신하는 식으로 insert를 진행한다.

 

self.cnt[i]는 i값이 바뀔 때마다 1씩 더해주면 되고, self.end[i]는 i가 마지막에 도착했을 때 1을 더해주면 되므로 위와 같이 구현하면 된다.

 

따라서 1차원으로 trie를 구현하면 종합적으로 아래와 같다.

class Trie:
    def __init__(self, alpha, leng, maxword, func):
        # alpha: the number of characters
        # max_node: the maximum number of nodes
        # func: function that translates characters to index num
        self.alpha = alpha
        self.maxnode = leng * maxword
        self.trie = [0] * (alpha * self.max_node)
        self.cnt = [0] * self.max_node
        self.end = [0] * self.max_node
        self.num = 0
        self.f = func

    def insert(self, word):
        i = 0
        self.cnt[i] += 1
        for ch in word:
            idx = self.f(ch)
            pt = self.alpha * i + idx
            if self.trie[pt] == 0:
                self.num += 1
                self.trie[pt] = self.num
            i = self.trie[pt]
            self.cnt[i] += 1
        self.end[i] += 1

 

이를 이용하면 메모리를 효율적으로 사용하면서 트라이 문제를 해결할 수 있다. 트라이는 문자열을 다루는 자료구조를 넘어서 이진수를 비트 단위로 트라이에 저장하는 아이디어를 활용해서 푸는 XOR 문제도 있을 정도로 범용성이 높은 자료구조이므로 1차원으로 구현하는 방법을 알아두면 상당히 편리할 것이다.

 

search prefix count나 search word 등의 함수도 구현해야 할 때가 있지만, 이는 cnt 또는 end 배열을 이용해서 들어가서 탐색하면 되므로 생략하겠다.

 

그러면 다음 글에서는 kmp를 확장하여 일대다 매칭이 가능한 아호 코라식을 살펴보겠다. 트라이 관련 문제를 몇 가지 추천하니 연습이 필요한 독자는 아래 문제들을 해결해보기를 바란다.

 

<BOJ 14725 개미굴> https://www.acmicpc.net/problem/14725: 트라이 기본 문제

<BOJ 13505 두 수 XOR> https://www.acmicpc.net/problem/13505: 이진수를 트라이를 이용해 관리하여 풀 수 있는 문제

<BOJ 28122 아이템> https://www.acmicpc.net/problem/28122: 트라이에서 dfs를 돌리는 문제