Week 3-2. [Trie] 208. Implement Trie (Prefix Tree) (Python3)

문제 해결 과정

✍️ 문제 요약

https://leetcode.com/problems/implement-trie-prefix-tree/?envType=study-plan-v2&envId=top-interview-150 

Implement Trie (Prefix Tree) - LeetCode

➡️ insert, search, startsWith 메서드를 가지는 Trie 자료 구조를 직접 구현하라.

 

🔎 어떤 식으로 접근했는가?

insert()

루트 노드부터 하나씩 자식 노드를 순회하면서 word의 구성 요소가 없을 때마다 새로운 노드를 추가한다.

 

search()

마찬가지로 루트 노드부터 자식 노드를 순회하는데, 이번에는 word의 구성 요소가 없을 때 바로 False를 반환한다.

 

startsWith()

마찬가지로 루트 노드부터 자식 노드를 순회하는데, 이번에는 prefix의 구성 요소가 없을 때 바로 False를 반환한다.

 

🎯 풀이

재귀를 이용한 풀이이다.

class Node:
    def __init__(self):
        self.children = dict()
        self.isEndOfWord = False

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

    def insert(self, word: str) -> None:
        self._insert(self.root, word)

    def _insert(self, node: Node, word: str) -> None:
        if len(word) == 0:
            node.isEndOfWord = True
            return
        c = word[0]
        child = node.children.get(c)
        if child is None:
            node.children[c] = Node() # 새로운 노드 추가
            child = node.children.get(c)
        self._insert(child, word[1:])

    def search(self, word: str) -> bool:
        return self._search(self.root, word)

    def _search(self, node: Node, word: str) -> bool:
        if len(word) == 0:
            return node.isEndOfWord
        c = word[0]
        child = node.children.get(c)
        if child is None: 
            return False
        return self._search(child, word[1:])

    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for x in prefix:
            child = node.children.get(x)
            if child == None: return False
            node = child
        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:
    node = self.root
    for c in word:
        if not c in node.children:
            node.children[c] = Node()
        node = node.children[c]
    node.isEndOfWord = True

위 코드의 시간 복잡도는 중위 순회를 통해 이진 트리의 모든 노드를 한 번씩 방문하므로 O(n)이다. 

문제에서 입력의 최대는 10^4이므로 감당할 수 있는 정도이다.

 

다른 사람의 코드를 탐구해보기

재귀를 사용하지 않아 더욱 쉽고 이해하기 좋다.

마찬가지로 시간 복잡도는 이진 트리의 모든 노드를 한 번씩 방문하므로 O(n)이다. 

 

class Node:
    def __init__(self):
        self.children = dict()
        self.isEndOfWord = False

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

    def insert(self, word: str) -> None:
        node = self.root
        for c in word:
            if c not in node.children:
                node.children[c] = Node()
            node = node.children[c]
        node.isEndOfWord = True

    def search(self, word: str) -> bool:
        node = self.root
        for c in word:
            if c in node.children:
                node = node.children[c]
            else: return False
        return node.isEndOfWord

    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for c in prefix:
            if c in node.children:
                node = node.children[c]
            else: return False
        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)