Trie 자료구조 [Python / 파이썬]

1. Trie 자료구조란?

※ 입력되는 문자열을 Tree 형식으로 만들어 진행되어 보다 빠르게 문자열 검색이 가능한 자료구조이다.

보통 list에서 문자열이 존재하는지 확인하기 위해서는 O(n) (n은 list의 길이)이라는 시간이 걸리는데

Trie 알고리즘을 사용하면 O(m) (m은 문자열의 길이) 이라는 짧은 시간이 소요되기 때문에 엄청 효율적이라고 할 수 있다.

2. 어떤 경우에 적합한가?

※문자열을 검색하는 문제에서 입력되는 문자열이 많을 경우 자주 사용된다.

빠른 시간복잡도 덕분에 검색엔진 사이트에서 제공하는 자동 완성 및 검색어 추천 기능에서 Trie 자료구조를 사용한다.

3. Trie 자료구조의 형태

1. Trie 알고리즘은 노드를 이용한 Tree 형태로 이루어져 있다.

2. 문자열의 끝을 알리는 flag가 존재한다.

3. 아래 예시는 "A", "to", "tea", "ted", "ten", "i", "in", "inn"를 키로 둔 트라이 이다.

출처 : https://commons.wikimedia.org/wiki/File:Trie_example.svg

4. 시간 복잡도

트라이는 다음과 같은 시간 복잡도를 갖는다.

제일 긴 문자열의 길이를 m, 총 문자열들의 수를 n이라 할 때

생성시 시간 복잡도 : O(m*n)
삽입 시간 복잡도 : O(m)
탐색시 시간 복잡도 : O(m)

5. 트라이의 구현 (Python)

5.1 Node 구현

class Node(object):
    def __init__(self, key, data=None):
        self.key = key
        self.data = data
        self.children = {}

노드에는 세가지가 필요하다.

key - 값으로 입력될 문자
data - 문자열의 종료를 알리는 flag. ( 이 예제 에서는 전체 문자열을 저장)
children - 자식노드를 저장

5.2 Trie 구현

class Trie:
    # 초기화 해드를 빈 노드로 설정
    def __init__(self):
        self.head = Node(None)

    # insert함수 - 트리를 생성하기 위한 함수
    def insert(self, string):
        # head노드부터 시작
        current_node = self.head
        
        # 문자열의 문자를 하나씩 확인
        for char in string:
            # 현재 노드의 자식중에 문자가 없다면
            if char not in current_node.children:
                # 자식 노드 추가
                current_node.children[char] = Node(char)
            # 자식 중에 문자가 있다면 current_node를 자식 노드로 변경
            current_node = current_node.children[char]
        
        # 문자열을 끝까지 탐색했다면 마지막 노드에 data추가
        current_node.data = string
    
    # Trie에서 string이 있는지 찾는 함수
    def search(self, string):
        # head노드부터 시작
        current_node = self.head
        
        #문자열의 문자를 하나씩 확인
        for char in string:
            # 만약 현재 노드의 자식노드중 문자에 해당하는 노드가 존재한다면
            if char in current_node.children:
                # current_node를 자식 노드로 변경
                current_node = current_node.children[char]
            # 현재 노드의 자식노드중 문자에 해당하는 노드가 없다면
            else:
                return False
        # 문자열 끝까지 탐색하여 마지막 노드에 데이터가 존재한다면
        if current_node.data:
            return True
        # 문자열 끝까지 탐색하여 마지막 노드에 데이터가 없다면
        else:
            return False

    # BFS 방식으로 깊이를 늘려가며 prefix로 시작하는 문자열을 저장한 배열 반환
    def starts_with(self, prefix):
        # head 노드부터 시작
        current_node = self.head
        #prefix로 시작하는 문자열을 저장할 리스트
        words = []

        # prefix의 처음부터 마지막까지 탐색
        for p in prefix:
            # 만약 현재 노드의 자식노드중 문자에 해당하는 노드가 존재한다면
            if p in current_node.children:
                # current_node를 자식 노드로 변경
                current_node = current_node.children[p]
            # 현재 노드의 자식노드중 문자에 해당하는 노드가 없다면 
            else:
                return None

        # prefix의 마지막 노드
        current_node = [current_node]
        # 다음 노드를 저장할 배열
        next_node = []
        while True:
            #현재 노드부터 탐색
            for node in current_node:
                #만약에 데이터가 있다면
                if node.data:
                    #word 배열에 추가
                    words.append(node.data)
                # 현재 노드의 자식 노드 전부를 next_node 배열에 추가
                # list extend 와 expend 정리
                next_node.extend(list(node.children.values()))
            
            # 만약 다음 노드가 존재한다면
            if len(next_node) != 0:
                # 다음 노드들을 현재 노드로 변경
                current_node = next_node
                #next_node는 다시 비워준다
                next_node = []
            # 다음 노드가 없다면 종료
            else:
                break

        return words

Trie 에는 여러가지 함수가 필요하다.

head를 빈노드로 설정
insert 함수 - 트리를 생성하기 위한 함수이다. 입력된 문자열의 문자를 하나씩 children에 확인 후 저장하고 문자열을 다 돌면 마지막 노드의 data에 문자열을 저장한다.
search 함수 - 문자열이 존재하는지에 대한 여부를 리턴하는 함수이다. 문자열을 하나씩 돌면서 확인 후 마지막 노드가 data가 존재한다면 True를, 그렇지 않거나 애초에 children에 존재하지 않는다면 False를 리턴한다.
starts_with 함수 - prefix단어로 시작하는 단어를 찾고 배열로 리턴하는 함수이다. prefix단어까지 tree를 순회 한 이후 그다음부터 data가 존재하는 것들만 배열에 저장하여 리턴한다.

Deeppago's study note