백준 2644번 : 촌수계산 in Python
이 문제는 여러 사람들에 대한 부모 자식 관계가 주어졌을 때 주어진 두 사람의 촌수를 계산하는 문제이다. 단, 각 사람의 부모는 최대 한 명만 주어진다.
이를 해결하기 위해서는 문제에서도 직관적으로 알 수 있듯이 각 사람들을 하나의 노드로 하고 부모 자식 관계는 간선으로 표현한 트리를 만들면 된다. 특히 문제의 조건 중에 각 사람의 부모가 최대 한 명이라는 조건에 의해 각 노드마다 부모 노드는 두 개 이상 될 수 없어 이를 구현하는데 문제가 없다.
트리는 각 노드마다 부모 노드의 번호를 저장하는 배열 parent를 만들어 구현한다. 예를 들어 n 9일 때 부모 자식 관계가 주어질 때 1 2라는 input이 주어지면 2의 부모가 1이라는 의미이므로 parent[2]의 값을 1로 한다. 이때 루트 노드, 즉 부모가 없는 노드를 표현하기 위해 초기값은 자기 자신의 노드 번호로 한다.
input으로 주어진 부모 자식 관계를 parent에 반영한 뒤 이제 주어진 두 사람의 촌수를 계산해야 한다. 촌수는 두 사람의 공통된 조상 노드 중 가장 높이가 낮은(깊이가 깊은) 노드를 찾아, 해당 조상 노드와 두 사람을 나타내는 노드 간의 촌수를 계산한 뒤 더해야 한다.
이를 위해 두 사람을 a, b라 할 때 a와 b의 조상 노드의 촌수를 계산하기 위한 배열을 cnt_a, cnt_b로 만들고, 초기값을 -1로 한다. 그리고 cnt_a[a]와 cnt_b[b]는 0으로 한다. 그 다음 a와 b의 모든 조상 노드에 대해 촌수를 계산해 cnt_a와 cnt_b에 반영하고, a와 b의 모든 공통 조상 중 촌수의 합이 가장 적은 값을 출력한다. 이는 a와 b의 공통 조상이 여러 명일 경우가 있을 수 있기 때문이다.
이를 코드로 구현하면 다음과 같다.
import sys
n = int(input())
a, b = map(int, input().split())
m = int(input())
# 부모 자식 관계를 트리로 나타내는 배열
parent = [i for i in range(n+1)]
for _ in range(m):
x, y = map(int, input().split())
parent[y] = x
# a와 b의 조상 노드의 촌수를 계산하기 위한 배열로 만약 조상 노드가 아닌 경우 -1이 저장될 것이다.
cnt_a = [-1 for _ in range(n+1)]
cnt_b = [-1 for _ in range(n+1)]
# 초기값을 설정한다.
cnt_a[a] = 0
cnt_b[b] = 0
temp_a = a
temp_b = b
# 결과값으로 공통 조상 노드의 촌수의 합의 최솟값을 구해야 하기에 초기값을 매우 큰 값으로 한다.
res = sys.maxsize
# 모든 조상 노드에 대해 촌수를 계산한다.
while parent[temp_a] != temp_a:
cnt_a[parent[temp_a]] = cnt_a[temp_a] + 1
temp_a = parent[temp_a]
while parent[temp_b] != temp_b:
cnt_b[parent[temp_b]] = cnt_b[temp_b] + 1
temp_b = parent[temp_b]
for i in range(1, n+1):
# i가 a와 b의 공통 조상인 경우
if cnt_a[i] != -1 and cnt_b[i] != -1:
res = min(res, cnt_a[i]+cnt_b[i])
if res != sys.maxsize:
print(res)
# a와 b의 공통 조상이 없는 경우
else: print(-1)