문제풀이/그래프

[Python/파이썬] 백준 11437번 LCA

딜레이레이 2024. 4. 10. 10:57
 

11437번: LCA

첫째 줄에 노드의 개수 N이 주어지고, 다음 N-1개 줄에는 트리 상에서 연결된 두 정점이 주어진다. 그 다음 줄에는 가장 가까운 공통 조상을 알고싶은 쌍의 개수 M이 주어지고, 다음 M개 줄에는 정

www.acmicpc.net

 

코드

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

n = int(input())
graph = [[] for _ in range(n+1)]    # 노드 간 연결
for _ in range(n-1):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

# 각 노드 레벨 구하기(BFS)
level = [-1]*(n+1)
parent = [i for i in range(n+1)]
q = deque([1])
level[1] = 0
while q:
    now = q.popleft()
    for nx in graph[now]:
        if level[nx] == -1:
            q.append(nx)
            level[nx] = level[now]+1
            parent[nx] = now


def find_lca(a, b):  # a, b의 LCA 찾기
    if level[a] < level[b]:  # a가 더 큰 레벨을 갖도록
        a, b = b, a
    # a, b 레벨 같게 맞추기
    while level[a] != level[b]:
        a = parent[a]
    # 같은 부모 찾기
    while a != b:
        a = parent[a]
        b = parent[b]
    return a


# 입력받은 두 정점의 가장 가까운 공통 조상(LCA) 출력
m = int(input())
for _ in range(m):
    a, b = map(int, input().split())
    print(find_lca(a, b))

 

가장 가까운 공통 조상을 구하는 알고리즘인 LCA(Lowest Common Ancestor)를 사용하면 된다.

 

1. 모든 노드의 레벨과 부모 노드를 구한다.

LCA 알고리즘을 사용하기 위해서는 각 노드의 레벨(깊이, depth)과 부모 노드를 알아야 한다. 이를 위해 루트 노드인 1번 노드부터 BFS로 탐색하여 level과 parent 리스트에 각 노드의 레벨과 부모 노드를 기록해준다.

 

2. LCA

입력으로 받은 a, b 두 노드의 가장 가까운 공통 조상(LCA)을 찾아주는 함수이다. 이 함수는 3단계로 나눌 수 있다.

 

① a가 b보다 큰 레벨을 갖도록 하기 위해 b의 레벨이 더 크다면 a와 b를 바꿔준다.

② a와 b의 레벨을 갖게 맞추기 위해 a의 레벨이 b와 같아질 때까지 a를 a의 부모로 바꿔주는 과정을 반복한다.

③ 이제 같은 레벨에 있는 a와 b를 a==b가 될 때까지 각각의 부모로 거슬러 올라가는 과정을 반복한다. 같아지면 그 때의 a(=b) 값을 리턴해준다.