[PS] BOJ 5052 / 전화번호 목록

[PS] BOJ 5052 / 전화번호 목록
Thumbnail: Photo by mdreza jalali (Unsplash)
문제 링크: https://www.acmicpc.net/problem/5052

풀이

​트라이 활용하기

트라이에 관한 자세한 설명은 이전 글을 참고해주세요.

이번 문제는 전화번호를 하나의 단어로 생각해 트라이에 저장하는 방법으로 해결할 수 있습니다.

전화번호의 숫자를 한 자리 씩 트라이의 노드로 저장하고, 마지막 자리에 해당하는 노드에는 전체 전화번호를 값으로 저장합니다. 한 전화번호부의 모든 전화번호를 트라이에 저장했다면, 트라이를 DFS로 탐색하며 말단 노드가 아닌 노드에 값이 저장되어있는지 확인하면 됩니다. 중간 노드에 값이 있다면 전화를 걸 때 이후 노드의 번호를 입력하지 못하고 바로 해당 번호로 전화가 걸리기 때문입니다.

class TrieNode:
    """Trie를 구성하는 노드 구현"""
    def __init__(self, key, value = None):
        self.key = key
        self.value = value
        self.children = {}
    
    @property
    def is_end(self):
        return len(self.children) == 0

    def add_child(self, key_idx, value):
        """이 노드에 새로운 자식 노드를 추가한다.
        루트 노드에서부터 재귀적으로 호출되며, 전체 전화번호를 한 자리씩 노드로 만들어 저장한다. 
        전화번호의 마지막 자리에 해당하는 노드에는 value 값도 저장한다.
        """
        key_str = value[key_idx]
        child = self.children.get(key_str, None)
        if child is None:
            child = TrieNode(key_str, None)
            self.children[key_str] = child
        
        if key_idx == len(value) - 1: # 전화번호의 마지막 자리 노드인 경우 value에 전화번호 저장
            child.value = value
        else: # 그렇지 않다면 재귀적으로 자식 노드를 만든다.
            child.add_child(key_idx + 1, value)

def check_trie(node):
    if node.is_end:
        return True
    
    res = True
    for child in node.children.values():
        if child.value is not None and not child.is_end:
            return False
        else:
            res = res and check_trie(child) 
    return res

Trie 및 전화번호부 검사 코드

전체 코드

input = open(0).readline

class TrieNode:
    """Trie를 구성하는 노드 구현"""
    def __init__(self, key, value = None):
        self.key = key
        self.value = value
        self.children = {}
    
    @property
    def is_end(self):
        return len(self.children) == 0

    def add_child(self, key_idx, value = None):
        key_str = value[key_idx]
        child = self.children.get(key_str, None)
        if child is None:
            child = TrieNode(key_str, None)
            self.children[key_str] = child
        
        if key_idx == len(value) - 1:
            child.value = value
        else:
            child.add_child(key_idx + 1, value)

def check_trie(node):
    if node.is_end:
        return True
    
    res = True
    for child in node.children.values():
        if child.value is not None and not child.is_end:
            return False
        else:
            res = res and check_trie(child) 
    return res

root = TrieNode(None)
for _ in range(int(input())):
    root.children.clear()
    N = int(input())
    for _ in range(N):
        root.add_child(0, input().rstrip())
    
    print("YES" if check_trie(root) else "NO")

solution.py