[백준 3584번] 가장 가까운 공통 조상 (C++, PyPy3)

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

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

이번 포스트에서는 백준 알고리즘 문제인 가장 가까운 공통 조상을 포스트하겠습니다.

 

해당 문제는 solved.ac 기준 골드 4에 위치한 문제입니다.

 

먼저 문제입니다.

입력으로 테스트 케이스 T와, 노드수 N, 간선 정보 N - 1개가 T번 주어집니다. 그리고 해당 케이스 마지막 입력으로 가장 가까운 공통 조상을 찾고자 하는 두 노드를 입력받습니다.

 

가장 가까운 공통 조상(LCA)은 두 노드의 부모 노드로 접근하면서 같은 노드에서 만날때 까지 반복하면 구할 수 있습니다.

 

이를 위해 먼저 setDepth() 함수와 LCA() 함수를 작성합니다.

 

setDepth() 함수는 파라미터로 노드 a와 부모 노드를 저장하는 parent 배열, 깊이를 저장하는 depth 배열을 받습니다.

parent[a]가 0 이라면 루트 노드이므로 그대로 리턴합니다.

그 외의 경우 새로운 노드 a가 들어왔으므로 부모 노드로 올라가며 재귀적으로 깊이를 1 증가시킵니다.

 

다음으로 LCA() 함수는 파라미터로 LCA를 구할 두 노드 x, y와 parent, depth를 받습니다.

공통 조상을 찾아야 하므로 두 노드의 깊이를 맞춰줍니다.

그리고 깊이가 같아졌다면 부모 노드로 올라가며 만날때 까지 반복하고 그 결과를 리턴합니다.

 

먼저 C++ 코드입니다.

 

#include <iostream>
 
using namespace std;
 
void setDepth(int a, int* parent, int* depth) {
    if(parent[a] == 0)
        return;
    else {
        setDepth(parent[a], parent, depth);
        depth[a] = depth[parent[a]] + 1;
    }
}
 
int LCA(int x, int y, int* parent, int* depth) {
    while(depth[x] > depth[y])
        x = parent[x];
 
    while(depth[x] < depth[y])
        y = parent[y];
 
    while(x != y) {
        x = parent[x];
        y = parent[y];
    }
 
    return x;
}
 
int main() {
    int t;
    cin >> t;
 
    int parent[10001] = {0,};
    int depth[10001];
 
    for(int i = 0; i < 10001; i++)
        depth[i] = 1;
 
    for(int i = 0; i < t; i++) {
        int n;
        cin >> n;
        
        for(int k = 0; k < n - 1; k++) {
            int a, b;
            cin >> a >> b;
            parent[b] = a;
        }
        
        for(int k = 1; k <= n; k++)
            setDepth(k, parent, depth);
 
        int t1, t2;
        cin >> t1 >> t2;
 
        cout << LCA(t1, t2, parent, depth) << "\n";
 
        for(int k = 0; k < 10001; k++) {
            parent[k] = 0;
            depth[k] = 1;
        }
    }
 
    return 0;
}

다음은 Python 코드입니다.

 

import sys
sys.setrecursionlimit(10**4+1)
input = sys.stdin.readline
 
def setDepth(a, parent, depth) :
    if parent[a] == 0 :
        return
    else :
        setDepth(parent[a], parent, depth)
        depth[a] = depth[parent[a]] + 1
 
def LCA(x, y, parent, depth) :
    while depth[x] > depth[y] :
        x = parent[x]
    while depth[x] < depth[y] :
        y = parent[y]
 
    while x != y :
        x = parent[x]
        y = parent[y]
    
    return x
 
parent = [0] * 10001
depth = [1] * 10001
 
t = int(input())
 
for i in range(t) :
    n = int(input())
 
    for k in range(n - 1) :
        a, b = map(int, input().split())
        parent[b] = a
 
    for k in range(1, n + 1) :
        setDepth(k, parent, depth)
 
    t1, t2 = map(int, input().split())
    print(LCA(t1, t2, parent, depth)
 
    for k in range(10001) :
        parent[k] = 0
        depth[k] = 1

해당 문제는 두 노드의 위치만 같게끔 맞춰주면 구하는데 어려움이 없는 문제였습니다.

단, 해당 파이썬 코드의 경우 PyPy3 에서만 가능하고 Python3 에서는 시간초과가 발생합니다.

따라서, 재귀로 깊이를 구하는 방식이 아닌 다른 방식으로 해결해야 합니다.

 

다음 포스트에서도 다른 백준 알고리즘 문제로 포스트 하도록 하겠습니다. 감사합니다.