[Python] LeetCode - 208. Implement Trie (Prefix Tree)

LeetCode 208번 Implement Trie (Prefix Tree) 문제는 문자열 탐색에 사용되는 Trie를 구현하는 문제다. 난이도는 Medium이다. 문자열 삽입, 탐색 메소드 및 추가적으로 startsWith 메소드를 구현해야 한다.

 

리트코드 208번 Implement Trie 문제 정보

알고리즘 분류

- Trie

난이도

- Medium

 

문제 요약

• Trie는 다진 트리 형태의 자료 구조이며, 효과적으로 문자열 데이터를 저장 및 탐색한다.
• Trie 클래스를 구현하라.

 

문제 풀이 과정

• 트리 형태로 구현해야 하기 때문에 먼저 트리의 노드가 되는 Node 클래스를 정의한다. 노드에 필요한 값은 character (value), 현재 노드에 해당하는 값이 존재하는지 여부 (data), 그리고 자식 노드 리스트(children)가 있다.
• 문자열 삽입 (insert)
  • 트리의 head부터 출발하여 입력할 문자열이 Trie에 존재하는지 확인하며 타고 내려간다. 필요한 노드가 없으면 추가한다.
  • 문자열이 Trie에 완성되면 노드에 해당 값이 존재함을 표시한다.
• 문자열 탐색 (search)
  • 트리의 head부터 출발하여 목표 문자열이 Trie에 존재하는지 확인하며 타고 내려간다.
  • Trie에 알맞은 노드가 존재하며, 노드에 해당 값이 존재하면 True를 반환한다.
• prefix 탐색 (startsWith)
  • 트리의 head부터 출발하여 목표 문자열이 Trie에 존재하는지 확인하며 타고 내려간다.
  • Trie에 알맞은 노드가 존재하면 True를 반환한다.

코드 - 1번 풀이

class Node:
    def __init__(self, val):
        self.value = val
        self.data = None
        self.children = []


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

    def insert(self, word: str) -> None:
        cur = self.head
        for s in word:
            is_existed = False
            for child_node in cur.children:
                if child_node.value == s:
                    is_existed = True
                    cur = child_node
                    break
            if not is_existed:
                new_node = Node(s)
                cur.children.append(new_node)
                cur = new_node
        
        cur.data = word

    def search(self, word: str) -> bool:
        cur = self.head
        for s in word:
            for child_node in cur.children:
                if child_node.data == word:
                    return True
                if child_node.value == s:
                    cur = child_node
                    break
        return False
        

    def startsWith(self, prefix: str) -> bool:
        cur = self.head
        for s in prefix:
            is_existed = False
            for child_node in cur.children:
                if child_node.value == s:
                    cur = child_node
                    is_existed = True
                    break
            
            if not is_existed:
                return False
        return True

• Trie 노드에 해당하는 character를 value에, 현재 노드까지의 문자열이 존재하는 경우 문자열을 data에, 자식 노드들을 children에 저장한 방식이다.
• 시간복잡도: insert - O(26*L), search - O(26*L), startsWith - O(26*L) (L: Trie에 저장할/탐색할 문자열의 길이)
  • children의 최대 길이가 알파벳 개수인 26이며, 저장 또는 탐색하려는 문자열의 길이만큼만 탐색하면 되기 때문이다.
• 공간복잡도: O(26 * 총 노드 수)

코드 - 2번 풀이

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


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

    def insert(self, word: str) -> None:
        cur = self.head
        for s in word:
            if s in cur.children.keys():
                cur = cur.children[s]
            else:
                new_node = Node()
                cur.children[s] = new_node
                cur = new_node
        
        cur.data = word

    def search(self, word: str) -> bool:
        cur = self.head
        for s in word:
            if s in cur.children.keys():
                if cur.children[s].data == word:
                    return True
                cur = cur.children[s]
        return False
        

    def startsWith(self, prefix: str) -> bool:
        cur = self.head
        for s in prefix:
            if s in cur.children.keys():
                cur = cur.children[s]
            else:
                return False
        return True

• 현재 노드까지의 문자열이 존재하는 경우 문자열을 data에 저장하고, 자식 노드들을 children에 { value: Node }의 Map 형태로 저장한 방식이다.
• 시간복잡도: insert - O(26*L), search - O(26*L), startsWith - O(26*L) (L: Trie에 저장할/탐색할 문자열의 길이)
  • children의 최대 길이가 알파벳 개수인 26이며, 저장 또는 탐색하려는 문자열의 길이만큼만 탐색하면 되기 때문이다.
• 공간복잡도: O(26 * 총 노드 수)