메뉴 닫기

Python3 Trie 자료구조 구현

class Node:
    def __init__(self):
        self.count = 0
        self.pointer = [None for i in range(26)]

class Trie:
    def __init__(self):
        self.root = Node()
        self.str_to_int = {}

        for i in range(97, 123):
            self.str_to_int[chr(i)] = i - 97
        
    def insert(self, string):
        self.root = self._insert(self.root, string)

    def _insert(self, node, string):
        if not node: #노드가 없다.
            node = Node()
            
        if not string: #종말 노드이다.
            node.count += 1
        else: #종말 노드가 아니다.
            node.pointer[self.str_to_int[string[0]]] = self._insert(node.pointer[self.str_to_int[string[0]]], string[1:])

        return node

    def find(self, string):
        node = self.root

        while node and string:
            node = node.pointer[self.str_to_int[string[0]]]
            string = string[1:]

        if not node:
            return 0
        if node.count == 0:
            return 0
        else:
            return node.count


trie = Trie()
trie.insert('aaaabb')
print(trie.find('aaaabb'))
print(trie.find('aaaab'))
print(trie.find('aab'))
print(trie.find('aaaabbbbb'))
trie.insert('xyz')
print(trie.find('xy'))
print(trie.find('yx'))
print(trie.find('xyz'))
trie.insert('xyz')
print(trie.find('xyz'))

문자열 검색에 자주 사용되는 자료구조인 Trie이다. 대학 수업에서는 따로 가르쳐 주지 않은 기억인데, 코딩 테스트에서는 꽤나 자주 이용되는 자료 구조인 것 같다. (2017, 2019 카카오 블라인드 테스트의 문제 하나 씩의 정해에서 사용되었다.)

이것보다 간단하게 구현하려면, 사전 자료구조를 계속 중첩하면서 구현하는 방법도 있다.

별도의 노드와 트라이 객체를 분리해서 작성할 필요 없이 노드 객체만을 사전과 카운트로 작성해주면 더 간단하게 구현이 가능하다.