Python/Algorithm

트라이

하효닝 2022. 2. 27. 04:43

트라이 (Trie)

  • 각 노드에 문자를 저장하는 트리 기반의 자료구로, 트리를 아래쪽으로 순회하면 단어 하나가 나온다.

  • 트라이는 길이가 K인 문자열이 주어졌을 때 O(K) 시간에 해당 문자열이 유효한 접두사인지 확인할 수 있다.
  • 먄악 트리에서 연관된 접두사를 반복적으로 검색해야 하는 경우, 트리의 현재 노드를 참조값으로 넘긴다면 매번 검색할 때마다 루트에서 시작할 필요없이 자식 여부만 확인하면 된다.

 

구현

# Node 클래스
# key: 해당 노드의 문자, child: 자식 노드
# data: 문자열이 끝나는 위치를 알려주는 역할 (해당 노드에서 끝나는 문자열이 없으면 None)
class Node:
    def __init__(self, key, data=None):
        self.key = key
        self.data = data
        self.child = dict()

# Trie 클래스
class Trie:
    def __init__(self):
        self.root = Node(None)

    # 문자열 삽입
    def insert(self, word):
        curr_node = self.root

        # 삽입할 word의 각각의 문자에 대해 자식 Node를 만든다.
        for char in word:
            # 자식 Node들 중 같은 문자가 없으면 새로운 Node 생성
            if char not in curr_node.child:
                curr_node.child[char] = Node(char)
            # 같은 문자가 있으면 노드를 따로 생성하지 않고 해당 노드로 이동
            curr_node = curr_node.child[char]

        # 문자열이 끝난 지점의 노드의 data 값에 해당 문자열 입력
        curr_node.data = word

    # 문자열 조회
    def search(self, word):
        # 가장 위에 있는 노드에서부터 탐색 시작
        curr_node = self.root

        for char in word:
            if char not in curr_node.child:
                return False
            curr_node = curr_node.child[char]

        # 탐색이 끝난 후 해당 노드의 data 값이 존재한다면 탐색 성공
        if curr_node.data != None:
            return True
        return False