[백준 파이썬] 11438 LCA 2

백준 11438

여기에 있는 강의 내용을 바탕으로 구현해봤다.

from collections import deque

N = int(input())
edges = [list(map(int, input().split())) for _ in range(N-1)]
M = int(input())
pairs = [list(map(int, input().split())) for _ in range(M)]

# 간선들을 인접 리스트로 표현
graph = [[] for _ in range(N+1)]
for i, j in edges:
    graph[i].append(j)
    graph[j].append(i)


# 너비 우선 탐색으로 노드들의 깊이와 부모 노드 탐색
def bfs():
    q = deque([1])
    parent = [-1 for _ in range(N+1)]
    parent[1] = 0
    depth = [-1 for _ in range(N+1)]
    depth[1] = 0

    while q:
        node = q.popleft()
        for next_node in graph[node]:
            if parent[next_node] == -1:
                parent[next_node] = node
                depth[next_node] = depth[node] + 1
                q.append(next_node)

    return parent, depth

# 너비 우선 탐색에서 얻은 정보로 2**k 단계씩 올라가며 부모 노드를 찾는 dp 함수
def dp(d):
    for k in range(1, K+1):
        for i in range(1, N+1):
            if d[k-1][i] == -1:
                d[k][i] = -1
            else:
                d[k][i] = d[k-1][d[k-1][i]]
    return d


parent, depth = bfs()
K = 0
while 2**K < max(depth):
    K += 1
    if 2**K > max(depth):
        K -= 1
        break

dp_parent = [[-1 for _ in range(N+1)] for _ in range(K+1)]
dp_parent[0] = parent

dp_depth = [[-1 for _ in range(N+1)] for _ in range(K+1)]
dp_depth[0] = depth

dp_parent = dp(dp_parent)
dp_depth = dp(dp_depth)


for i, j in pairs:
	# 깊이가 더 깊은 노드가 트리의 리프에 더 가깝다 -> low
    if dp_depth[0][i] > dp_depth[0][j]:
        high, low = j, i
    else:
        high, low = i, j

    # 낮은 위치(리프에 가까운) 노드를 루트에 가까운 노드와 깊이를 맞춰준다.
    k = K # 2**k가 트리의 전체 높이보다 커질 수 없으므로 높이보다 낮은 가장 큰 2의 거듭제곱 수터
    # 아래의 노드의 깊이가 더 깊다면
    while dp_depth[0][high] < dp_depth[0][low]:
    	# 아래 노드의 깊이에서 2**k 만큼 올라갔을 때 위 노드보다 높아진다면
        if dp_depth[0][low] - 2**k < dp_depth[0][high]:
        	# k-1을 해줘서 더 위 노드보다 낮은 깊이가 되도록 k 감소
            k -= 1
        else:
        	# 위 노드보다 높이가 낮거나 같다면 아래 노드를 2**k 단계 위의 노드로 변경
            low = dp_parent[k][low]

    # 공통 노드 찾기
    k = K
    # k가 0이 될 때까지
    while k >= 0:
    	# 둘의 노드가 같다면 -> 최소가 아닌 공통 조상일 경우도 있다
        if dp_parent[k][high] == dp_parent[k][low]:
        	# 같은 노드를 가르키지 않을 때까지 k 감소
            k -= 1
        else:
        	# 서로 다른 노드를 가르킨다면 각 노드를 부모 노드로 업데이트
            high = dp_parent[k][high]
            low = dp_parent[k][low]

	# k=0일 때는 부모 노드가 같은 노드를 가르킬 경우와 현재 노드가 같은 노드일 경우이다
    if high == low:
        print(high)
    else:
        print(dp_parent[0][high])