[Python] 백준 11437 - LCA

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

문제

N(2 ≤ N ≤ 50,000)개의 정점으로 이루어진 트리가 주어진다. 트리의 각 정점은 1번부터 N번까지 번호가 매겨져 있으며, 루트는 1번이다.

두 노드의 쌍 M(1 ≤ M ≤ 10,000)개가 주어졌을 때, 두 노드의 가장 가까운 공통 조상이 몇 번인지 출력한다.

입력

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

출력

M개의 줄에 차례대로 입력받은 두 정점의 가장 가까운 공통 조상을 출력한다.

 

---

풀이

LCA는 어떤 트리에서 두 정점이 주어질 때 정점들의 최소 공통 조상을 찾는 문제이며, 아주아주아주 트리에서 기본적이고 필수적인 개념이다.

그런데도 나는 이 문제를 D5가 될때까지도 안풀고 있었다. 이제 슬슬 트리도 해야하는데....

 

사실 1년 전쯤에 이 문제를 풀다가 포기한 적이 있었다. 그래서 Solved AC 프로필에 한동안 "메모: 11437 풀기"라는 글이 적혀있었다.

이제 PS를 꽤 잘하는 나의 머리로 이 문제를 이제서야 쉽게 푼다.

 

LCA를 구하는 방법은 여러가지가 있지만, 가장 쉽고 직관적인 방법은 "직접 위로 한칸씩 같아질 때 까지 올려보는 것"이다.

 

그렇게 하기 위해서, 맨 처음에 DFS로 트리의 형태를 잡아주고, 이 때 각 노드의 조상인 parent 배열과 깊이인 level 배열을 채워준다.

 

그리고, 맨 처음에 주어진 두 정점의 레벨이 다르다면 같아질 때까지 맞춰준다.

그리고, 두 정점을 같아질 때까지 한 칸씩 올리는 것이다.

 

코드의 구현은 아래와 같다.

 

import sys
input = sys.stdin.readline
N = int(input())

graph = [[] for _ in range(N+1)]
for i in range(N-1):
    u, v = map(int, input().split())
    graph[u].append(v)
    graph[v].append(u)

visited = [False]*(N+1)
visited[1] = True
parent = [i for i in range(N+1)]
level = [0]*(N+1)
stack = [1]
while stack:
    now = stack.pop()
    for i in graph[now]:
        if not visited[i]:
            visited[i] = True
            parent[i] = now
            level[i] = level[now] + 1
            stack.append(i)

#print(parent)
#print(level)

M = int(input())
for i in range(M):
    a, b = map(int, input().split())
    if a == b:
        print(a)
        continue

    if level[a] > level[b]:
        while level[a] > level[b]:
            a = parent[a]
    elif level[b] > level[a]:
        while level[b] > level[a]:
            b = parent[b]
    if a == b:
        print(a)
        continue
    while parent[a] != parent[b]:
        a = parent[a]
        b = parent[b]
    print(parent[a])

 

 

참고로, 고수들은 DFS를 재귀 안쓰고 stack으로 푼다. 이게 훨씬 빠르다.

(BFS에서 큐 대신 스택쓰는거랑 똑같다)

이 글을 읽는 여러분들도 이제 DFS에서 재귀를 쓰는 것을 졸업해보도록 하자.