[BOJ] 11437 - LCA (Python)

 

 

아이디어

 

최소 공통 조상 알고리즘을 이용해 가장 가까운 공통 조상을 알아낸다.

 

 

풀이

 

너비 우선 탐색 함수 bfs에서는 트리를 탐색하며 현재 노드의 깊이, 부모 노드를 저장하도록 한다.

최소 공통 조상을 찾는 함수 find는 두 노드의 깊이를 일정하게 맞춘 뒤, 공통 조상을 찾아 반환한다.

딕셔너리 memo에는 이미 조상을 찾았던 두 노드의 최소 공통 조상을 저장해두어 find 함수를 재호출하지 않도록 한다.

 

입력 받은 간선들을 2차원 리스트 graph에 추가한 후, 함수 bfs를 호출하여 노드 깊이와 부모 노드 정보를 저장한다.

M개의 두 노드에 대한 최소 공통 조상을 찾아 출력하고, 메모이제이션을 통해 실행 시간을 단축한다.

 

import sys

input = sys.stdin.readline

def bfs():
    queue = [1]
    visited[1] = 1
    while queue:
        idx = queue.pop(0)
        for i in graph[idx]:
            if not visited[i]:
                depth[i] = depth[idx] + 1
                parent[i] = idx
                queue.append(i)
                visited[i] = 1

def find(a, b):
    if depth[a] > depth[b]:
        a, b = b, a
    while depth[b] > depth[a]:
        b = parent[b]
    while a != b:
        a = parent[a]
        b = parent[b]
    return a

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)
parent = [0]*(N+1)
depth = [0]*(N+1)
visited = [0]*(N+1)
bfs()
M = int(input())
memo = dict()
for _ in range(M):
    a, b = map(int, input().split())
    if not memo.get((a, b), 0):
        memo[(a, b)] = memo[(b, a)] = find(a, b)
    print(memo.get((a, b)))

 

 

메모이제이션을 적용하지 않으면 Python 3로는 시간 초과가 나고, PyPy3로 겨우 통과할 수 있었다.

# 더 느린 코드
M = int(input())
for _ in range(M):
    a, b = map(int, input().split())
    print(find(a, b))