[백준] 11437번 : LCA (Python)

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

 

import sys
from collections import deque

input = sys.stdin.readline

N = int(input()) # N의 개수
tree = [[] for _ in range(N+1)] # 인접 리스트
parent = [0]*(N+1) # 부모 노드 저장
depth = [0]*(N+1) # 깊이 저장
result = {}

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

# 너비 우선 탐색
def BFS(start):
    q = deque([start])
    depth[start] = 1 # 최초 노드에 depth 1 초기화

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

# 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)]

    # 깊이 맞추기
    while(depth[a] != depth[b]):
    
        if depth[a] > depth[b]:
            a = parent[a] # 부모 노드로 올리기
        else:
            b = parent[b]

    # LCA 찾기
    while(a != b):
        a, b = parent[a], parent[b]

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

BFS(1) # BFS 실행
M = int(input()) # 질의의 개수 

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

 

 

[문제풀이]

---

 

두 노드의 가장 가까운 공통 조상 → LCA

 

대놓고 LCA 알고리즘을 사용하라고 낸 문제다.

노드 개수가 50,000개라 해도 트리로 만들면 높이가 얼마 안된다.

또, 질의도 작다 보니 비교적 데이터가 크지 않아 일반적인 LCA로 구해도 충분히 구할 수 있는 문제!

 

1. 인접리스트에 트리 저장
2. DFS, 또는 BFS 알고리즘을 사용하여 배열 초기화(depth 구하고, 부모 노드 저장)
3. 두 노드의 깊이 맞추기
4. 최소 공통 조상 찾기

 

일반적인 최소 공통 조상은 위의 단계를 거쳐 구한다.

이때 문제 조건 상 연결된 두 노드->방향 없는 트리기 때문에, 인접리스트에 양방향으로 저장해야 한다.

 

여기서 단계에서 주의해야 할 특이점이 파이썬은 아무래도 느리기 때문에....

딕셔너리에 구한 최소 공통 조상을 저장해놓고, 다음 번에 또 같은 노드로 질의가 들어올 때 바로 return 해주는 코드를 추가해야 한다.

if (a, b) in result:
   return result[(a, b)]

 

바로 이 부분!

이걸 메모이제이션이라 한다.

 

이걸 하지 않으면 계속 시간 초과가 난다ㅠㅠ