[ BOJ ] 3584 : 가장 가까운 공통 조상 ( GOLD 4 ) / Python

>> 문제 바로가기 (https://www.acmicpc.net/problem/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 이하의 정수로 이름 붙여집니다.

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

출력

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

풀이 과정

먼저 루트 노드를 찾고, 찾고자 하는 두 노드를 찾기 위해 루트 노드부터 탐색을 시작하도록 하였다.

만약 찾고자 하는 노드에 도착했다면 해당 노드 번호를 반환하지만, 자식 노드가 없는 리프 노드에서는 -1을 반환하였다.

 

부모 노드에서 -1 정보만 받고 아무 노드 번호에 대한 정보를 받지 못했다면, 해당 부분 트리에서는 찾고자 하는 노드를 포함하지 않는다고 판단하고 -1을 반환했다.

만약 찾고자 하는 노드에 대한 정보를 1개 받았다면, 더 상위 높이에서 공통 조상을 찾아야 한다고 판단하고 자식 노드에서 받은 노드 번호를 반환했다.

찾고자 하는 노드에 대한 정보를 2개 받았다면, 해당 노드가 가장 가까운 공통 조상이라고 판단하고 해당 노드 번호를 반환했다.

 

이 기능을 구현한 함수를 수행하면, 재귀를 통해 가장 가까운 공통 조상을 반환받을 수 있다.

import sys
input = sys.stdin.readline

t = int(input().rstrip())
for _ in range(t):
    n = int(input().rstrip())
    tree = [[] for _ in range(n+1)]
    parent = [-1 for _ in range(n+1)]

    for _ in range(n-1):
        a, b = map(int, input().rstrip().split())
        tree[a].append(b)
        parent[b] = a
    
    parent[0] = 0
    root = parent.index(-1)

    order = list(map(int, input().rstrip().split()))
    
    def findAnscestor(now):
        if now in order:
            return now
        
        if len(tree[now]) == 0:
            return -1
        
        child_value = 0
        child_cnt = 0
        for next in tree[now]:
            result = findAnscestor(next)
            if result != -1:
                child_value = result
                child_cnt += 1
        
        if child_cnt == 0: 
            return -1
        elif child_cnt == 1:
            return child_value
        else:
            return now
        

    print(findAnscestor(root))