[PS] BOJ 5052 / 전화번호 목록
문제 링크: 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 resTrie 및 전화번호부 검사 코드
전체 코드
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