[백준] 3584. 가장 가까운 공통 조상 - Python

[Gold IV]

 

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

3584번: 가장 가까운 공통 조상

 

풀이

처음 보는 LCA (Lowest Common Ancestor) 문제.

DFS를 통해 Root 노드부터 각 노드의 깊이(depth)를 저장하고,

입력된 두 노드의 depth를 맞춘 후, 공통 조상이 나올 때까지 그대로 거슬러 올라가는 식으로 해결하였다.

 

이 풀이의 경우 O(N)의 Time Complexity로 해결한 경우이며,

개선된 LCA 알고리즘의 경우 O(logN) 의 복잡도로 해결할 수 있다고 한다.

LCA 알고리즘에 대해서는 여러 곳에서 잘 설명해 놓았으니, 이 곳에는 따로 적지 않겠다.

 

AC.

import sys
sys.setrecursionlimit(int(1e5))

T = int(input())

for _ in range(T):
    N = int(input())
    parent = [0] * (N+1)
    depth = [0] * (N+1)
    cal = [False] * (N+1)
    tree = [[] for _ in range(N + 1)]

    for _ in range(N-1):
        a, b = map(int, input().split())
        tree[a].append(b)
        tree[b].append(a)
        parent[b] = a

    # 루트 노드 찾기 (parent가 0(초기값)이면 root)
    root = 0
    for i in range(1, N+1):
        if parent[i] == 0:
            root = i
            
    # 루트 노드부터 depth 저장.
    def dfs(x, dep):
        cal[x] = True
        depth[x] = dep
        for y in tree[x]:
            if cal[y]:
                continue
            parent[y] = x
            dfs(y, dep+1)
    
    # Lowest Common Ancestor
    def LCA(x, y):
        while depth[x] != depth[y]:
            if depth[x] > depth[y]:
                x = parent[x]
            else:
                y = parent[y]
        
        while x != y:
            x = parent[x]
            y = parent[y]
        
        return x

    # root부터 dfs를 통해 depth 저장
    dfs(root, 0)
    
    # LCA
    n, m = map(int, input().split())
    print(LCA(n, m))