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

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

 

코드

import sys
        
def parent(a):
    p = []
    p.append(a)  # 자기 자신도 조상으로 포함
    for _ in range(N):
        if tree[a] == []:
            break
        else:
            p.append(tree[a])
        a = tree[a]    
    return p
        
T = int(input())
for _ in range(T):
    N = int(input())
    tree = [[] for _ in range(N+1)]
    a_par = []
    b_par = []
    for _ in range(N-1):
        A,B = map(int,sys.stdin.readline().split())
        tree[B] = A   # 자신의 부모노드 저장
    a,b = map(int,sys.stdin.readline().split())
    a_par = parent(a)
    b_par = parent(b)
    common = [x for x in a_par if x in b_par]
    print(common[0])

 

설명

먼저 tree 리스트에 자신의 부모노드에 해당하는 번호를 추가해줬다.

그리고 나서 입력받은 두 노드의 조상들을 모두 찾아야 되기 때문에 parent 함수를 사용하였다. 입력받은 노드부터 시작해서 계속 거슬러 올라갔고 만약 부모 노드가 [ ]가 되면 루트 노드까지 도착한 것이기 때문에 for문을 빠져나간다. 이때 자기 자신도 조상으로 추가해야 한다.  

예제를 예로 들면  tree = [[ ], 3, [ ], 2, 3, 1]와 같이 저장이 될텐데 0번째는 제외하고 인덱스가 2일 때 [ ]이므로 2가 루트 노드임을 알 수 있다. 

두 노드의 조상들을 모두 구한 뒤 두 리스트의 공통 수를 common 리스트에 뽑아낸다. 이때 순서는 입력 노드와 가까운 순으로 되어있기 때문에 common 리스트에서 가장 처음에 있는 값을 출력하면 가장 가까운 공통 조상을 찾을 수 있다.