Trie (트라이)

Python/파이썬 자료구조 알고리듬

Trie (트라이)

컴닥 2022. 12. 5. 11:04

[참고] tistory 코드의 가독성을 높이는 법

나무 위키의 간결한 설명.

https://namu.wiki/w/%ED%8A%B8%EB%9D%BC%EC%9D%B4

트라이 - 나무위키

이 저작물은 CC BY-NC-SA 2.0 KR에 따라 이용할 수 있습니다. (단, 라이선스가 명시된 일부 문서 및 삽화 제외) 기여하신 문서의 저작권은 각 기여자에게 있으며, 각 기여자는 기여하신 부분의 저작권

namu.wiki

설명이 너무 잘 되어 있다. 코드까지 ㄷㄷㄷ

https://ko.wikipedia.org/wiki/%ED%8A%B8%EB%9D%BC%EC%9D%B4_(%EC%BB%B4%ED%93%A8%ED%8C%85)

트라이 (컴퓨팅) - 위키백과, 우리 모두의 백과사전

위키백과, 우리 모두의 백과사전. "A", "to", "tea", "ted", "ten", "i", "in", "inn"를 키로 둔 트라이. 이 예제에는 모든 자식 노드가 알파벳 순으로 왼쪽에서 오른쪽으로 정렬되어 있지는 않다. (루트 노드

ko.wikipedia.org

간단하게는 딕셔너리로 작성할 수 있다.

def insert(root, word):
    current = root
    for letter in word:
        if letter not in current:
            current[letter] = {}
        current = current[letter]  # 자식 딕셔너리로 내려간다. 핵심적인 부분.
    current['*'] = word  # 단어의 끝을 '*'으로 마킹함.


# def insert(root, word):
#     current = root
#     for letter in word:
#         current.setdefault(letter, {})  # 파이써닉하게 멋을 낸다면... 
#         current = current[letter]
#     current['*'] = word


def search(root, word):  # insert 함수와 비슷한 구조. 
    current = root
    for letter in word:
        if letter not in current:
            return False
        current = current[letter]  
    if current['*'] == word:
        return True
    return False


def solve(words, queries):
    trie = {}
    for word in words:
        insert(trie, word)
    print(trie)
    for query in queries:
        print(query, search(trie, query))


solve(['abc', 'bcd', 'efg', 'abcd'], ['abcd', 'abd', 'aefg', 'efg'])

클래스를 이용한다면 일단은 이렇게 작성할 수 있고....

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

    def insert(self, word):
        current = self.root
        for letter in word:
            if letter not in current:
                current[letter] = {}
            current = current[letter]
        current['*'] = word

    def search(self, word):
        current = self.root
        for letter in word:
            if letter not in current:
                return False
            current = current[letter]
        if current['*'] == word:
            return True
        return False


def solve(words, queries):
    trie = Trie()
    for word in words:
        trie.insert(word)
    print(trie.root)
    for query in queries:
        print(query, trie.search(query))


solve(['abc', 'bcd', 'efg', 'abcd'], ['abcd', 'abd', 'aefg', 'efg'])

기능의 확장이 필요하다면 Node 클래스를 만들어 보자.

몇 개의 단어가 이 노드를 지나가는 지 확인할 수 있도록 카운터를 붙여 보겠다. (알고리듬 문제 풀이에 많이 사용됨)
단어의 끝을 확인하는 인스턴스 변수 end도 같이 ...

class Node:
    def __init__(self):
        self.count: int = 0
        self.children: dict = {}
        self.end: bool = False

    def __repr__(self):
        return f'{self.count} {self.end} {self.children}'


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

    def insert(self, word: str) -> None:
        current = self.root
        current.count += 1
        for letter in word:
            if letter not in current.children:
                current.children[letter] = Node()
            current = current.children[letter]
            current.count += 1
        current.end = True

    def search(self, word: str) -> (bool, int):
        current = self.root
        for letter in word:
            if letter not in current.children:
                return False, current.count
            current = current.children[letter]
        if current.end:
            return True, current.count
        return False, current.count
    
    # 가독성 나쁨
    # def search(self, word: str) -> (bool, int):
    #     current = self.root
    #     for letter in word:
    #         if letter not in current.children:
    #             break
    #         current = current.children[letter]
    #     else:
    #         return current.end, current.count
    #     return False, current.count


def solve(words: list[str], queries: list[str]) -> None:
    trie = Trie()
    for word in words:
        trie.insert(word)
    print(trie.root)
    for query in queries:
        print(query, trie.search(query))


solve(['abc', 'bcd', 'efg', 'abcd'], ['abcd', 'abd', 'aefg', 'efg'])

Trie (트라이)

관련 문항