LeetCode 211 - Design Add and Search Words Data Structure

단어를 추가하고 검색하는 자료구조를 설계하는 Medium 문제다. 일반 문자만 검색한다면 Trie를 그대로 따라가면 되지만, 마침표는 어떤 문자와도 일치하는 wildcard라서 해당 위치의 모든 자식 노드를 확인해야 한다. 그래서 저장은 Trie, 검색은 DFS로 구성했다.

문제 링크 & 설명

문제 링크: 211. Design Add and Search Words Data Structure
요약: addWord로 단어를 저장하고, search로 완전히 일치하는 단어가 있는지 확인하는 WordDictionary를 구현한다. 검색어의 .은 임의의 영문 소문자 하나와 일치한다.

각 Trie 노드는 알파벳 26개에 대응하는 자식 배열과, 현재 위치에서 단어가 끝나는지를 나타내는 is_end를 가진다. 접두사가 존재하는 것과 완성된 단어가 저장된 것은 다르기 때문에 종료 표시가 따로 필요하다.

접근 방법

addWord는 단어의 문자를 순서대로 읽으며 자식 노드가 없을 때 생성한다. 마지막 노드에서는 is_end를 true로 바꾼다.

for ch in word:
    idx = ord(ch) - ord('a')
    if not cur.children[idx]:
        cur.children[idx] = TrieNode()
    cur = cur.children[idx]

cur.is_end = True

검색은 현재 노드와 문자열 깊이를 DFS 상태로 사용한다. 일반 문자는 해당 인덱스의 자식 하나만 따라가고, .을 만나면 존재하는 모든 자식에서 다음 깊이를 탐색한다. 그중 하나라도 남은 문자열과 일치하면 true다.

if ch == '.':
    for child in node.children:
        if self.dfs(child, depth + 1, word):
            return True

문자열 끝까지 도착했을 때는 노드의 존재만 확인하면 안 된다. app만 저장된 상태에서 ap를 검색하면 경로는 존재하지만 단어는 저장되지 않았기 때문에, is_end를 반환해야 한다.

복잡도 분석

단어 추가: 길이를 L이라고 할 때 시간 복잡도는 O(L)이고, 새 노드가 모두 필요하면 추가 공간도 O(L)이다.
검색: wildcard가 없으면 O(L)이다. .이 여러 개 있으면 각 위치에서 최대 26개 자식으로 분기하므로 최악의 경우 Trie의 여러 노드를 탐색한다. 재귀 호출 스택은 단어 길이만큼 쌓여 O(L) 공간을 사용한다.
전체 저장 공간: 저장된 Trie 노드 수에 비례한다.

구현 코드

class TrieNode:
    def __init__(self):
        self.children = [None] * 26
        self.is_end = False

class WordDictionary:
    def __init__(self):
        self.root = TrieNode()

    def addWord(self, word: str) -> None:
        cur = self.root

        for ch in word:
            idx = ord(ch) - ord('a')
            if not cur.children[idx]:
                cur.children[idx] = TrieNode()
            cur = cur.children[idx]

        cur.is_end = True

    def dfs(self, node, depth, word) -> bool:
        if not node:
            return False
        if len(word) == depth:
            return node.is_end

        ch = word[depth]
        if ch == '.':
            for child in node.children:
                if self.dfs(child, depth + 1, word):
                    return True
        else:
            idx = ord(ch) - ord('a')
            next_node = node.children[idx]
            if next_node and self.dfs(next_node, depth + 1, word):
                return True

        return False

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

요약 및 회고

Trie의 기본 검색은 경로 하나를 따라가는 연산이지만 wildcard 하나가 들어오면 탐색 문제가 된다. 같은 자료구조라도 검색 조건에 따라 단순 순회가 DFS로 바뀌는 지점이 이 문제의 핵심이었다. is_end가 접두사와 완성된 단어를 구분한다는 점도 Trie 구현에서 빠뜨리면 안 되는 조건이었다.