백준 2644번: 촌수계산 (트리)

백준 2644번: 촌수계산

부모 중 하나만 표시하는 촌수를 그리면 트리다.
트리를 파이썬으로 어떻게 구현하더라 고민하다 보니
subtree의 level을 이용해서 촌수를 구하는 것이 편할 것 같았다.

두 사람이 친척이라면 조상을 타고 올라가면서 -1이 나오기 전에 마지막 조상(root)은 무조건 같은 사람이다.
그리고 맨 위 조상까지 가기 전에 subtree의 root가 있을 수도 있다.

A와 B의 촌수를 구할 때
A와 B의 root로부터 depth 차이는 2다.

p라는 자신의 부모를 저장하는 리스트가 있다.
2만큼 부모를 찾아가면 B와 세대가 같은 A의 조상 p[p[a]를 찾을 수 있다.
이 때까지 촌수는 2다.

p[p[a]]와 B가 한 번씩 부모를 확인하며 겹치는 조상을 찾는다.
한 세대를 올라갈 때마다 촌수는 2씩 증가한다.
조상이 같을 때 subtree의 root를 발견하고 촌수 3을 구할 수 있다.

def get_level(x):
    level = 0
    while parent[x] != -1:	# root까지
        x = parent[x]
        level += 1
    return level

def get_ancestor(x, depth):
    for _ in range(depth):
        x = parent[x]
    return x

def get_chon(xp, yp, chon):
    while xp != yp or xp == -1:
        if xp == -1 or yp == -1:	# 친척 X
            return -1
        
        xp = parent[xp]
        yp = parent[yp]
        chon += 2
    return chon

n = int(input())
parent = [-1] * (n + 1)

x, y = map(int, input().split())
m = int(input())
for i in range(m):
    x_, y_ = map(int, input().split())
    parent[y_] = x_

depth = get_level(x) - get_level(y)

if depth > 0:
    x = get_ancestor(x, depth)

elif depth < 0:     # B가 A보다 아래 세대일 때
    depth = -depth
    y = get_ancestor(y, depth)

print(get_chon(x, y, depth))