[백준] 11437 - LCA (Python)

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

 

워냑 유명한 유형이다보니 (나도 풀어보지는 않았지만 이 유형에 대해서 들어봤었다.) 정형화된 알고리즘이 있을 것 같은데, 우선 스스로의 고민으로 먼저 풀어보기로 했다.

 

어떤 트리가 주어질 때, 임의의 두 노드에 대해 그 두 노드의 최소 공통 조상 노드를 찾는 것이 문제의 요구사항이다.

이때 트리의 루트는 1번 노드이다.

(나는 이 조건을 못 보고 어떤 루트로 잡든 성립하는 문제인가보다~ 하고 1번 노드를 루트로 잡고 풀었는데, 생각해보면 두 노드 중에 한 노드를 매번 루트로 잡으면 그 노드가 최소 공통 조상이 되어버리니 루트는 정해져있어야 의미있는 문제였다.)

 

처음에는 특정 노드에서부터 루트를 향해 올라가면서 공통 조상 노드를 찾는 방법을 떠올렸다.

그런데 왼쪽에 25000, 오른쪽에 25000 개 노드가 존재하는 두 갈래 트리를 생각할 때, 양 끝에서 루트노드에서부터 공통 조상 노드를 찾는다고 하면 총 25000번의 비교를 해야 한다.

문제에서 주어지는 쿼리의 횟수는 최대 1만번 이므로 1만 * 25000 = 2억 5천만 번의 연산을 해야 한다.

(지금 생각해보니 이 풀이도 3초라는 시간 안에 충분히 풀 수 있을 것 같다.)

 

하지만 나는 이렇게 풀면 시간이 너무 오래 걸릴 것 같아서 다른 풀이를 고민했다.

그 결과 내가 생각한 아이디어는 매개변수 탐색을 활용한 풀이이다.

트리를 쭉 그렸을 때 각 depth 별 노드 정보를 알고 있다고 해보자.

그러면 depth 0 에 대해서는 루트노드만 하나 존재할 것이고, 그 밑으로 노드들이 쭉쭉 존재할 것이다.

 

이때 어떤 임의의 두 노드의 최소 공통조상은 반드시 존재한다.

depth 0 의 루트 노드는 어떤 노드 쌍에 대해서도 공통 조상이기 때문이다.

 

만약 어떤 두 노드의 LCA 가 k depth 를 갖는 노드라고 해보자.

그러면 k-1 depth 를 갖는 그 노드의 조상도, k-2 depth 를 갖는 그 노드의 조상도 모두 '두 노드의 공통 조상' 이 될 수 있다.

즉, depth 0 부터 공통 조상이 될 수 있는지 여부를 쭉 나열하면 0부터 k 까지는 모두 가능하다가 k+1 부터 공통조상이 되지 않는다.

그렇다면 이 문제는 (k, k+1) 경계를 찾는 매개변수 탐색 문제로 바뀐다.

 

특정 depth 에 어떤 두 노드의 공통 조상이 존재하는지 판별하려면 어떻게 해야할까?

나는 나이브하게 최초 DFS 탐색을 돌리면서 각 depth 를 기준으로 노드 정보를 저장했다.

그리고 특정 depth 에 대해서 두 노드의 공통 조상이 존재하는지 판별하기 위해 해당 depth 를 갖는 모든 노드에 대해 그 노드가 두 노드의 조상 노드인지 판별했다.

 

어떤 두 노드가 서로 부모-자식 관계를 갖는지 안 갖는지는 어떻게 판별할까?

DFS 탐색의 특성을 사용하면 어떤 노드를 방문한 시간과 그 노드를 빠져나와서 되돌아온 시간을 timestamp 로 찍을 수 있다.

 

부모 노드의 timestamp가 (a, b) 자식 노드의 timestamp 가 (c, d) 라면 이 둘 사이에는 반드시 (a < c < d < b) 관계가 성립한다.

만약 이 관계가 성립하지 않으면 형제 노드 (sibling) 이거나 부모,자식 관계가 역전되어있다.

 

따라서 어차피 depth 판별을 위해서 DFS를 돌릴 때 timestamp 를 같이 찍고, 이 정보를 활용하여 부모-자식 관계를 판별해주면 된다.

 

정리하면 다음과 같다.

 

1. left = 0, right = 두 노드 중 더 낮은 depth + 1 로 초기값을 잡는다. (left 는 반드시 가능한 후보, right 는 반드시 불가능한 후보)

2. mid = (left + right) / 2 로 후보 depth 를 정한다.

3. 후보 depth 에 대해 해당 depth 를 갖는 모든 노드를 순회하면서 두 노드와 각각 부모-자식 관계인지 확인한다.

4. 두 노드와 모두 부모 자식 관계를 갖는 노드가 해당 depth 에 있다면, 그 노드를 정답의 후보로 하고 left = mid 로 설정한다.

5. 두 노드와 모두 부모 자식 관계를 갖는 노드가 해당 depth 에 없다면, right = mid 로 설정한다.

 

이분 탐색이 끝나면 정답의 후보를 출력한다.

import sys
sys.setrecursionlimit(5*(10**4)+100)
input = sys.stdin.readline


def DFS(node, start, now_depth):
    end = start
    next_start = start + 1
    depth[now_depth].append(node)
    node_depth[node] = now_depth
    for next_node in graph[node]:
        if not visit[next_node]:
            visit[next_node] = True
            end = DFS(next_node, next_start, now_depth+1)
            next_start = end + 1

    timestamp[node] = (start, end+1)
    return end+1


def is_parent(parent, child):
    if parent == child:
        return parent
    parent_start, parent_end = timestamp[parent]
    child_start, child_end = timestamp[child]
    return parent_start < child_start and child_end < parent_end


def possible(now_depth, node1, node2):
    for check_node in depth[now_depth]:
        if is_parent(check_node, node1) and is_parent(check_node, node2):
            return check_node

    return -1


n = int(input())
graph = [[] for _ in range(n)]
visit = [False]*n
depth = [[] for _ in range(n+1)]
node_depth = [0]*n
for _ in range(n-1):
    s, e = map(int, input().split())
    graph[s-1].append(e-1)
    graph[e-1].append(s-1)

timestamp = [None]*n
visit[0] = True
end = DFS(0, 0, 0)

m = int(input())
for _ in range(m):
    n1, n2 = map(int, input().split())
    n1, n2 = n1-1, n2-1

    n1_depth = node_depth[n1]
    n2_depth = node_depth[n2]

    s, e = 0, min(n1_depth, n2_depth) + 1
    ans = 0
    
    while s + 1 < e:
        now_depth = (s+e) // 2
        result = possible(now_depth, n1, n2)
        if result != -1:
            ans = result
            s = now_depth
        else:
            e = now_depth

    print(ans+1)

 

자료구조 시간에 배웠던 DFS 의 timestamp 특성을 활용하게 되어 재미있게 풀었다.

 

---

 

이제 정해를 찾아보았다.

정해는 크게 2가지 풀이가 있다.

 

1. O(Md) 풀이

내가 처음에 생각한 그 풀이이다.

찾으려고 하는 두 노드의 깊이가 다르다면 먼저 깊이를 맞춰주고, 그 다음에 한칸씩 위로 올라가면서 같은 부모 노드를 만날 때까지 반복한다.

최악의 경우, 깊이가 50000이 될 수 있으므로 50000 깊이 노드와 0번 깊이 노드의 공통 조상을 위 알고리즘으로 찾으면 5억번의 연산을 해야한다. 내가 푼 LCA 문제의 경우 이 풀이도 통과된다.

 

2. O(log d) 풀이

1번 풀이를 조금 더 효율적으로 풀기 위한 시도로 깊이를 맞춰주고 부모를 찾을 때 선형적으로 하나씩 올라가면서 찾는 것이 아니라 이분탐색으로 자신 기준 k 번째 부모를 빠르게 찾는다. 이때 자신의 부모 정보를 저장할 때 parent[i][j] 를 i 번 노드의 2^j 번째 부모 노드를 저장하는 아이디어를 사용한다. (그냥 j 번째 부모 노드를 저장하면 저장 공간이 5만 x 5만개가 필요할 것이다)

그러면 이분탐색으로 k번째 부모를 찾기 위해서 j 값을 선별하는 과정을 log 횟수로 줄일 수 있다.

 

이 부분은 나중에 코드로 직접 구현해보고 LCA 2 문제를 풀이해보겠다.