[Python/파이썬] 백준 3584번 가장 가까운 공통 조상

3584번: 가장 가까운 공통 조상

 

코드

from collections import deque
import sys
input = sys.stdin.readline

for _ in range(int(input())):
    n = int(input())
    parent = [-1]*(n+1)
    children = [[] for _ in range(n+1)]
    for _ in range(n-1):
        a, b = map(int, input().split())
        parent[b] = a
        children[a].append(b)
    x, y = map(int, input().split())

    # 루트 노드 찾기
    root_node = -1
    for i in range(1, n+1):
        if parent[i] == -1:
            root_node = i
            break
            
    # 각 노드의 depth 구하기
    q = deque([(root_node, 0)])
    depth = [0] * (n+1)
    while q:
        now, d = q.popleft()
        for nx in children[now]:
            q.append((nx, d+1))
            depth[nx] = d+1
            
    # 최소 공통 조상 (LCA)
    if depth[x] > depth[y]:
        x, y = y, x
    while depth[x] < depth[y]:  # x, y 깊이 같게 맞춰주기
        y = parent[y]
    for i in range(n):
        if x == y:
            print(x)
            break
        x = parent[x]
        y = parent[y]

 

어제와 같이 최소 공통 조상(LCA) 알고리즘을 활용한 문제이다. 위 코드는 크게 4가지 부분으로 나눌 수 있다.

 

1. 입력

간선 정보를 입력받을 때, 각 노드의 부모와 자식들 정보를 각각 저장한다.

 

2. 트리의 루트 노드 찾기

부모 노드의 값을 저장해놓은 parent 배열의 값이 초기값인 -1인 노드는 부모 노드가 없다는 뜻이므로 그 노드가 루트 노드이다.

 

3. 각 노드의 깊이 구하기

LCA 알고리즘을 수행하기 위해서는 각 노드의 깊이를 구할 필요가 있다. 그러므로 BFS 탐색을 이용하여 모든 노드를 탐색하고 깊이 정보를 depth 배열에 저장한다.

 

4. 최소 공통 조상 찾기

• 더 깊이가 깊은 노드가 y에 오도록 한다.
• x, y의 깊이가 같은 수준이 되도록 맞춰준다.
• 높이가 같아진 x, y에서 한 레벨씩 위로 올라오며 공통 조상, 즉 같은 노드 나올때까지 반복한다.