알고리즘

[python3] Trie 구현

https://github.com/Dev-Guccin 2022. 3. 19. 12:25

트라이는 문자열 탐색에 굉장히 유리한 알고리즘이다.

각 문자열에 대해서 트라이 형태로 자료구조를 만들면 이후에 문자열이 존재하는지 파악할때 빠르게 파악가능하다.

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


class Trie(object):
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):  # 문자열이 삽입한다. 단순 삽입
        curNode = self.head

        for char in string:
            if char not in curNode.children:
                curNode.children[char] = Node(char)  # 없는 경우 새롭게 추가
            # 같은 문자가 있으면 노드를 따로 생성하지 않고, 해당 노드로 이동
            curNode = curNode.children[char]
        # 문자열이 끝난 지점의 노드의 data값에 해당 문자열 입력
        curNode.data = string

    def find(self, string):  # Node의 data를 찾는과정
        curNode = self.head

        for char in string:
            if char in curNode.children:
                curNode = curNode.children[char]
            else:
                return False

        # 탐색이 끝난 후 노드의 data값이 존재한다.
        if curNode.data != None:
            return True


trie = Trie()
words = ["abcd", "abbdd", "front"]
for word in words:
    trie.insert(word)

print(trie.find("abcd"))