가장 가까운 공통 조상 [백준 3584 파이썬]

문제

트가 있는 트리(rooted tree)가 주어지고, 그 트리 상의 두 정점이 주어질 때 그들의 가장 가까운 공통 조상(Nearest Common Anscestor)은 다음과 같이 정의됩니다.

두 노드의 가장 가까운 공통 조상은, 두 노드를 모두 자손으로 가지면서 깊이가 가장 깊은(즉 두 노드에 가장 가까운) 노드를 말합니다.

예를 들어 15와 11를 모두 자손으로 갖는 노드는 4와 8이 있지만, 그 중 깊이가 가장 깊은(15와 11에 가장 가까운) 노드는 4 이므로 가장 가까운 공통 조상은 4가 됩니다.

루트가 있는 트리가 주어지고, 두 노드가 주어질 때 그 두 노드의 가장 가까운 공통 조상을 찾는 프로그램을 작성하세요

입력
첫 줄에 테스트 케이스의 개수 T가 주어집니다.

각 테스트 케이스마다, 첫째 줄에 트리를 구성하는 노드의 수 N이 주어집니다. (2 ≤ N ≤ 10,000)

그리고 그 다음 N-1개의 줄에 트리를 구성하는 간선 정보가 주어집니다. 한 간선 당 한 줄에 두 개의 숫자 A B 가 순서대로 주어지는데, 이는 A가 B의 부모라는 뜻입니다. (당연히 정점이 N개인 트리는 항상 N-1개의 간선으로 이루어집니다!) A와 B는 1 이상 N 이하의 정수로 이름 붙여집니다.

테스트 케이스의 마지막 줄에 가장 가까운 공통 조상을 구할 두 노드가 주어집니다.

출력
각 테스트 케이스 별로, 첫 줄에 입력에서 주어진 두 노드의 가장 가까운 공통 조상을 출력합니다.

전체 코드

import sys

sys.setrecursionlimit(10 ** 5)


# 첫 번째 노드에 대해 dfs를 해본다. 그러면 루트까지의 경로가 나올 것이고, 해당 경로를 list에 담아준다.
def dfs_one(tree: list[int], node: int, result: list[int]) -> list[int]:
    if tree[node] == 0: # 루트 노드라면, 루트 노드를 담고 리턴.
        result.append(node)
        return result
    else:
        result.append(node)
        return dfs_one(tree, tree[node], result)


# 첫 번째 dfs에서 얻었던 경로 정보를 바탕으로, 다른 노드에서의 경로 중 같은 것이 있는지 점검하며 진행한다.
def dfs_other(tree: list[int], node: int, one_result_set: set) -> int:
    # 루트 끝 까지 도달했을 경우 루트를 리턴
    if tree[node] == 0:
        return node
    else:
        # 굳이 set을 사용한 이유는, set의 경우 in 연산자가 O(1)이기 때문이다.
        # 만약 현재 node가 첫 번째 dfs에서 발견되었다면, 더 이상 dfs를 진행할 필요가 없다.
        if node in one_result_set:
            return node
        # 아직 발견되지 않았다면, 계속 진행한다.
        else:
            return dfs_other(tree, tree[node], one_result_set)


# 공통 조상 찾기 함수
def nearest_common_ancestor(tree: list[int], one: int, other: int) -> int:
    one_result = dfs_one(tree, one, [])
    nearest_ancestor = dfs_other(tree, other, set(one_result))
    return nearest_ancestor


test_case = int(input())
for i in range(test_case):
    n = int(input())
    nodes = [0] * (n + 1)
    for j in range(n - 1):
        parent, child = list(map(int, input().split()))
        nodes[child] = parent
    a, b = list(map(int, input().split()))
    ancestor = nearest_common_ancestor(nodes, a, b)
    print(ancestor)

개선점

시간 복잡도를 비교했을 때, 반복문 풀이보다 10배 이상 느렸다. 반복문이 당연히 재귀보다 빠르지만, 처음 푸는 유형이라 미처 생각을 못했다.
다음에는 재귀 대신에 반복문 방법으로 풀자.