[SWEA] 1248. 공통 조상

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV15PTkqAPYCFAYD 

SW Expert Academy

트리 자료를 저장하는 방법에 따라 쉽게 풀릴 수도, 아닐 수도 있는 문제였다.

나는 처음에 딕셔너리로 부모 : [자식1, 자식2] 을 받아줬는데 공통 조상을 찾을 때 코드 구현이 난감했다.

떠오른 것은 리스트에 [부모, 자식1, 자식2] 으로 받아주는 것과 클래스를 만들어주는 건데 리스트가 더 편할 것 같아서 이렇게 구현했다.

부모 찾기, 가장 가까운 공통 부모 찾기, sub_tree 개수 찾기 모두 함수를 따로 만들어주었다.

D5인데도 불구하고 정답률이 높길래 의아했는데 트리 구성 아이디어만 있으면 쉽게 풀리는 문제다.

 

def find_parent(node):
    parent_list = []
    while True:
        if tree[node][0] == 0 :
            return parent_list
        parent = tree[node][0]
        parent_list.append(parent)
        node = parent


def find_near_node():
    for p1 in node1_parent:
        for p2 in node2_parent:
            if p1 == p2:
                return p1


def find_subtree(node):
    stack = [node]
    cnt = 0
    while stack:
        now_node = stack.pop()
        cnt += 1
        if tree[now_node][1] != 0 :
            stack.append(tree[now_node][1])
        if tree[now_node][2] != 0 :
            stack.append(tree[now_node][2])
    return cnt


for tc in range(1, int(input())+1):
    v, e, node1, node2 = map(int, input().split())
    # 부모, 자식1, 자식2
    tree = [[0, 0, 0] for _ in range(v+1)]
    temp = list(map(int, input().split()))
    for i in range(0, e*2, 2):
        # 부모 - 자식
        tree[temp[i+1]][0] = temp[i]
        if tree[temp[i]][1] == 0:
            tree[temp[i]][1] = temp[i+1]
        else:
            tree[temp[i]][2] = temp[i+1]
    node1_parent = find_parent(node1)
    node2_parent = find_parent(node2)
    near_node = find_near_node()
    print('#{} {} {}'.format(tc, near_node, find_subtree(near_node)))