[백준] 11438번 : LCA2 (Python)

문제 : https://www.acmicpc.net/problem/11438

 

import sys
from collections import deque
from math import log2

input = sys.stdin.readline

N = int(input()) # N의 개수
tree = [[] for _ in range(N+1)] # 인접 리스트
parent = [0]*(N+1) # 부모 노드 저장
depth = [0]*(N+1) # 깊이 저장
visited = [False]*(N+1) # 방문 여부 배열
result = {} # 정답 쌍 저장

# 트리 저장
for _ in range(N-1):
    S, E = map(int, input().split()) 
    tree[S].append(E) 
    tree[E].append(S) # 인접리스트에 트리 저장

# 너비 우선 탐색
def BFS(start):
    q = deque([start])

    while q:
        now = q.popleft()
        visited[now] = True # 방문 설정
        
        for i in tree[now]:
            if not visited[i]: # 미방문 노드일때
                parent[i] = now # 부모 노드 저장
                depth[i] = depth[now]+1 # 깊이 저장
                q.append(i) # 자식 노드 저장
                
BFS(1) # BFS 실행

# Parent 배열 구하기
depth_max = max(depth) # 트리의 높이
K = int(log2(depth_max)) # 최대 가능 높이
P = [[0 for _ in range(N+1)] for _ in range(K+1)] # 부모 노드 배열
P[0] = parent # 바로 위 부모 노드 parent 배열에 저장

for i in range(1, K+1):
    for j in range(1, N+1):
        P[i][j] = P[i-1][P[i-1][j]] # 점화식을 이용해 Parent 배열 구함

# LCA 찾기
def LCA(S, E):

    # 항상 (작은 값, 큰 값)의 순서로 결과를 저장
    if S > E:
        S, E = E, S

    a, b = S, E # 함수 내에서 사용할 변수에 복사

    # 메모이제이션, 이미 정답 쌍에 저장된 값일 시 바로 반환
    if (a, b) in result:
        return result[(a, b)]

    # 깊이를 고려해 a, b swap
    if depth[a] > depth[b]:
        a, b = b, a
            
    # 깊이 맞추기        
    while(depth[a] != depth[b]):

        diff_depth = abs(depth[b] - depth[a]) # 두 노드의 depth 차이
        k = int(log2(diff_depth)) # 최대 k  
        b = P[k][b] # 부모 노드로 올리기

    # 두 노드가 같은지 check
    if(a == b):
        return a

    # LCA 찾기
    for k in range(K, -1, -1):

        if(P[k][a]!=P[k][b]): # 두 노드의 부모가 달라지는 값을 찾는다.
            
            a, b = P[k][a], P[k][b] # 찾을 시 저장            

    # 다른 노드라면 바로 위의 부모 노드가 LCA
    if(a!= b):
        a, b = P[0][a], P[0][b]

    result[(S, E)] = a # 메모이제이션을 위한 정답 저장
    return a # 찾은 최소 공통 조상 return

M = int(input()) # LCA를 구하고 싶은 M의 개수    

for i in range(M):
    S, E = map(int, input().split()) 
    print(LCA(S, E)) # LCA 찾기

 

 

 

[문제풀이]

---

 

노드의 개수(N)와 질의의 개수(M)이 매우 크다.

일반적인 최소 공통 조상 구하기 방식으로 구현하면 시간 초과가 발생하기 때문에, LCA 빠르게 구하기로 풀어야 한다.

 

1. 인접 리스트로 트리 데이터 구현

2. 탐색 알고리즘(DFS, BFS)을 이용해 각 노드의 깊이와, 바로 위 부모 노드를 구한다.

3. 점화식을 이용해 Parent 배열(부모 노드 배열)을 구한다

4. Parent 배열을 이용해 2^k만큼 빠르게 이동해 깊이를 맞춘다.

5. 깊이가 같아진 두 노드를 Parent 배열을 이용해 2^k만큼 이동해 최소 공통 조상을 찾는다.

k는 depth의 최댓값에서 1씩 감소한다.

 

LCA 빠르게 구하기 알고리즘이다.

 

여기서 k, 즉 몇번째 부모노드까지 저장할 지 구하는 건 위 코드로 구했다.

depth_max = max(depth) # 트리의 높이
K = int(log2(depth_max)) # 최대 가능 높이

 

k는 2^k <= 트리 높이여야 한다.

이걸 만족하는 k값을 찾고자 log2 함수를 사용했다.

 

두 노드의 높이를 맞춰줄때도 log2 함수를 사용했다.

diff_depth = abs(depth[b] - depth[a]) # 두 노드의 depth 차이
k = int(log2(diff_depth)) # 최대 k  

 

2^k <= 두 노드의 depth 차

이게 최대가 되는 k값을 구하기 위해 log 함수를 사용!

이렇게 구한 k값을 사용해서 한번에 이동해서 높이를 맞췄다.