[BOJ] 14725번 - 개미굴 (Trie 자료구조 )

이번에 포스팅 하게 된 것은 바로 문자열 알고리즘 중 하나인 Trie 자료구조 이다. 

맨 처음에 뭔지 모르고 있다가 한 번 알게 된 뒤로 신기하고 기억해두면 좋을 것 같아 한 번 글로써 정리 해보았다.

간단하게 말하자면 문자열 검색시 시간복잡도를 문자열의 길이인 M 즉 O(M)만큼 줄여주는 자료구조이다.

예를들어 여러개의 단어들이 있다고 가정하자. 

["ABC", "ABCD", "ABD", "ACBD", "ACD"] 의 단어들이 있다고 보겠다.

 

트라이 자료구조에서는 이 문자열들의 저장 형태는 아래의 그림처럼 저장이 된다. (하위 코드도 참조)

먼저 첫번째 원소인 "ABC" 를 대입 해보자.

"ABC" 대입

 

파이썬의 경우에는 dictionary를 이용하여 구현 할 수 있게 되고, C언어의 경우 포인터를 활용해서 구현 해놓은 블로그들을 보았다.

두번째 원소를 넣게 된다면 어떻게 될까?

 

"ABCD" 대입

이런식으로 세번째, 네번째 원소를 대입해보자.

 

"ABD", "ACBD" 대입

마지막 원소 대입.

 

최종 원소 대입

자, 이렇게 자료구조가 다 완성 되었다.

이렇게 자료구조를 넣는 시간은 총 O(MN) N: 문자열 개수 M: 문자열길이 가 될 수 있겠다.

 

만약 이런 자료구조 없이 다음 질문이 들어왔을 때 시간 복잡도는 얼마가 될 것인지 생각해보자.

 

질문 : "ABDE" 라는 문자열이 포함되어 있는가?

 

위의 자료구조 없이 만약 이 질문을 통해 답을 알아내기 위해서는 O(MN) 즉 각 문자열의 문자 하나하나를 비교해야 할 것이다.

즉, 하나를 찾는데 걸리는 시간만 O(MN)이라는 것이다.

 

하지만 트라이 자료구조를 활용했을 때, A->B->D 를 타고 내려오는 순간 D 에서 E를 찾을 수 없기 때문에

해당하는 단어가 자료구조에 없음을 나타내고 있다. 시간복잡도는 문자열의 길이인 O(M)이 된다.

 

간단하게 자료구조를 소개해 보았다. 아래의 코드를 참고하여 디버깅을 해보면 알아서 자료구조가 눈에 들어 올 것이다.

글쓴이 또한 처음 봤을 때 이해하지 못하고 직접 코드를 치면서 깨달았다.

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

    def add(self, word):
        tr = self.trie
        for w in word:
            if w not in tr:
                tr[w] = {}
            tr = tr[w]
        tr['*'] = word

    def search(self, word):
        tr = self.trie
        for w in word:
            if w in tr:
                tr = tr[w]
            else:
                return False
        if tr['*'] == word:
            return True

 

www.acmicpc.net/problem/14725

14725번: 개미굴

 

1. Trie 자료구조를 통해 개미들의 식량을 저장하여 본다.

 

2. 자료구조 안을 preorder로 탐색 한다.

 

위의 자료구조를 몇 번 연습해보고 이와 같은 문제를 풀어본다면 이해도가 훨씬 높아질 것이라고 생각한다.

글쓴이도 python 알고리즘 문제들을 풀면서 클래스를 써본 것은 처음이다. 앞으로 자료구조와 관련된 이런 문제들이 나오면

클래스로 구현하여 풀 예정이다.

 

from sys import stdin

std = stdin.readline

N = int(std())


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

    def add(self, data):
        tr = self.trie
        for d in data:
            if d not in tr:
                tr[d] = {}
            tr = tr[d]
        #tr['*'] = '*'


trie = Trie()
for i in range(N):
    info = list(std().split())[1:]
    trie.add(info)


def dfs(t, depth):
    if len(t) == 0:
        return

    t = dict(sorted(t.items()))
    # print(t)
    for key in t:
        for i in range(depth):
            print("--", end="")
        print(key)
        dfs(t[key], depth+1)


dfs(trie.trie, 0)