56. Implement Trie (Prefix Tree)

리트코드

• 트라이의 insert, search, startsWith 메소드를 구현하라.

1. 딕셔너리를 이용해 간결한 트라이 구현

class TrieNode:
    def __init__(self):
        self.word = False
        self.children = collections.defaultdict(TrieNode)

class Trie:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = TrieNode()
        

    def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        node = self.root
        for char in word:
            node = node.children[char]
        node.word = True
        
    
    # 단어 존재 여부 판별
    def search(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]
        return node.word
        
    #문자열로 시작 단어 존재 여부 판별
    def startsWith(self, prefix: str) -> bool:
        """
        Returns if there is any word in the trie that starts with the given prefix.
        """
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        return True
        


# Your Trie object will be instantiated and called as such:
# obj = Trie()
# obj.insert(word)
# param_2 = obj.search(word)
# param_3 = obj.startsWith(prefix)

• 트라이를 직접 구현해보는 문제다. 여기서는 딕셔너리를 이용해 가급적 간결한 형태로 풀이해본다.

• 먼저 트라이를 저장할 노드는 별도 클래스로 선언한다.

  • 메소드를 포함해 같은 클래스로 묶을 수도 있지만, 간결하게 구현하기 위해 별도로 구분해 선언한다.

• 트라이 연산을 구현할 별도 클래스를 선언하고 삽입 메소드를 구현한다.

 def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        node = self.root
        for char in word:
           if char not in node.children:
               node.children[char] = TrieNode()
            node = node.children[char]
        node.word = True

• 처음 Trie 클래스를 생성하게 되면 루트 노드로 별도 선언한 TrieNode 클래스를 갖게 되고, 삽입 시 루트부터 자식 노드가 점점 깊어지면서 문자 단위의 다진 트리(m-ary Tree) 형태가 된다.
• 만약 입력값이 apple인 경우, 삽입 코드는 다음과 같다.

t = Trie()
t.insert('apple')

<apple이 추가된 트라이>

• 이 그림에서 트라이는 다음 문자를 키로하는 자식 노드 형태로 점점 깊어지면서, 각각의 노드는 word 값을 갖는다.

  • 이 값은 단어가 모두 완성되었을 때만 True가 된다.

  • 즉 apple의 경우 단어가 모두 완성되는 e에 이르러서야 True 로 셋팅된다.

• 만약 appear, appeal 같은 문자가 추가로 삽입된다면 코드는 다음과 같다.

t.insert('appear')
t.insert('appeal')

<apple 트라이에 appear, appeal 단어가 추가된 트라이>

• 이 그림에서 트라이는, 같은 문자가 같은 자식을 타고 내려가다가, 달라지는 문자부터 서로 다른 노드로 분기된다.

• 마지막에는 appeal과 appear가 완성되는 l과 r노드가 각각 True로 셋팅된다.

• 여기까지가 삽입의 기본 원리다.

• 값이 존재하는지 확인하려면 어떻게 해야할까

  • search()는 단어가 존재하는지 여부를 판별하는 것이고

  • startsWith()는 해당 문자열로 시작하는 단어가 존재하는지 여부를 판별한다.

  • 둘 다 동일하게 문자 단위로 계속 깊이 탐색을 하고 search()의 경우에만 마지막에 word가 True인지 확인하면 될 것이다.

• search()를 구현해보자.

  • 문자열에서 문자를 하나씩 for 반복문으로 순회하면서 자식 노드로 계속 타고 내려간다. 그리고 마지막에 node.word여부를 리턴한다.

  • 단어가 완선된 트라이라면 True로 되어 있을 것이고, 이때 True가 결과로 리턴된다.

• startsWith()를 구현해보자.

  • search()와 거의 동일하나, 차이점은 node.word를 확인하지 않고, 자식 노드가 존재하는지 여부만 판별한다는 점이다.

  • 자식 노드가 존재한다면 for 반복문을 끝까지 빠져나올 것이고, True를 리턴한다.

• 코드를 줄여보자.

  • self.children을 defaultdict()로 선언한다면 insert()삽입 메소드에서 매번 if로 체크하는 구문은 없앨 수 있을 것 같다.