[Algorithm] 트라이(Trie)와 관련 문제들 (Python)

트라이란?

문자열을 효과적으로 탐색하기 위한 자료구조로, 문자열을 순차탐색하는 방식이 아닌 트리로 접근하기 위해 만들어진 자료구조입니다. 주로 Prefix Tree, digital search Tre, retrieval tree라고도 불리는데, retrieval tree의 trie를 따와 트라이라고 일반적으로 불립니다.

트라이는 문자열을 key로 사용하는 동적은 set 또는 연관배열을 저장하는 트리의 확장된 구조입니다. 문자열을 탐색할 때 단순하게 하나씩 비교하지 않고 key로 노드들을 탐색해나가기때문에 매우 효율적이라는 장점이 있습니다. 단, 빠르게 탐색 가능하지만 각 노드에서 자식들에 대한 포인터들을 배열로 모두 저장하고 있다는 점에서 저장 공간의 크기가 크다는 단점이 있습니다.

위 사진을 보시면 아시다싶이 루트 노드에서 하나씩 뻗어가면서 자식 노드에는 하나의 문자(key)가 들어있고 그 밑의 자식 노드드들은 바로 위 부모 노드와 공통의 접두사를 공유하고 있습니다.

트라이는 탐색과 삽입을 지원하는 노드로 구성된 트리여서, 탐색은 키 문자열과 연관된 값을 반환하고 삽입은 문자열과 그 값을 삽입합니다. 키의 길이가 m일때 탐색과 삽입 모두 $O(m)$ 시간에 수행됩니다.

구현 방법

관련 문제 : LeetCode - 208. Implement Tree(Prefix tree)

Implement Trie (Prefix Tree) - LeetCode

위 문제는 한번 꼭 보시는 것을 추천합니다! 트라이를 어떻게 구현할 수 있을지에 관한 문제거든요 ㅎㅎ 파이썬으로 Trie 클래스를 만들어봅시다.

class Trie:

    def __init__(self):
        self.trie = {}  # root 

    def insert(self, word):
        t = self.trie
        for c in word:
            if c not in t:
                t[c] = {}
            t = t[c]
        t['*'] = True # 끝이라는걸 알림

    def search(self, str)
        t = self.trie
        for c in word:
            if c not in t:
                return False
            t = t[c]
        return '*' in t

    def startsWith(self, prefix)
        t = self.trie
        for c in prefix:
            if c not in t:
                return False
            t = t[c]
        return True

한 문자를 입력할 때 해당 문자가 끝인 것을 알리기 위해 * 라는 키를 넣어줍니다. 탐색하는 함수의 경우에도 t = t[c] 가 자식노드를 계속해서 탐색해준다는 뜻입니다. 

 

Trie 문제

Leetcode 648. Replace Words

Replace Words - LeetCode

문제 풀이

class Solution:
    def replaceWords(self, dictionary: List[str], sentence: str) -> str:
        # 트라이 자료구조 만들기
        root = {}
        for dic in dictionary:
            t = root
            for c in dic:
                if c not in t:
                    t[c] = {}
                t = t[c]
            t['*'] = True

        answer = []
        for word in sentence.split():
            t = root
            res = ''

            if len(word) == 1:
                if word in t:
                    answer.append(word)
                    continue 

            flag = 0
            for c in word:
                if '*' in t:
                    answer.append(res)
                    flag = 1
                    break

                if c not in t:
                    answer.append(word)
                    flag = 1
                    break

                t = t[c]
                res += c

            if flag == 0:
                answer.append(res)

        return " ".join(answer)

먼저 트라이 구조를 만들어준 후 같은 prefix를 공유하는 단어가 트라이에 있으면 그 단어로 치환해주는 코드입니다. 있는지 없는지 여부는 *로 확인해주면 되고 검색하는 문자열보다 트라이에 있는 값이 더 짧은 경우를 확인하기 위해 flag를 같이 넣어주었어요.

사실 $O(N^3)$으로 풀 수 있는 문제고 통과는 되지만 효율성 측면에서 trie를 사용하는것이 더욱 좋다고 합니다 ㅎ.ㅎ 트라이 연습하실때 좋은 코드인 것 같아요.

---

다른 언어에서는 잘 모르겠는데, 파이썬에서는 아무래도 key-value로 들어가는 것이다 보니 딕셔너리 자료구조를 많이 사용하는 것 같습니다. 그냥 쉽게 코딩하려면 딕셔너리를 쓰고, 그렇지 않은 경우에는 따로 Node 클래스를 구현해주더라구요. 그래도 시간복잡도 측면에서는 비슷합니다.

트라이는 시간복잡도가 순차 탐색보다 훨씬 효과적이라는 장점이 있지만, 공간복잡도 측면에서는 별로 좋지가 않습니다. 당장 알파벳으로만 봐도.. a-z로 depth가 조금 깊어지면 ...... 엄청나게 많은 공간을 차지하게되기 때문에 코딩 환경에 따라 잘 확인해볼 필요가 있을 것 같아요 ㅎ.ㅎ