Trie (트라이)

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

 

나무 위키의 간결한 설명. 

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

트라이 - 나무위키

 

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

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

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

 

 

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

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'])

 

관련 문항

https://comdoc.tistory.com/entry/3%EC%B0%A8-%EC%9E%90%EB%8F%99%EC%99%84%EC%84%B1

[3차] 자동완성

https://comdoc.tistory.com/entry/2020-KAKAO-BLIND-RECRUITMENT%EA%B0%80%EC%82%AC-%EA%B2%80%EC%83%89

2020 KAKAO BLIND RECRUITMENT 가사 검색