백준 | 트리의 부모 찾기

백준 트리의 부모 찾기

문제

루트 없는 트리가 주어질 때, 각 노드의 부모를 구하는 프로그램을 작성하는 문제이다.

추가 조건은 다음과 같다.

[1] 트리의 루트는 1이라고 정의한다.
[2] 노드의 개수 N(2 ≤ N ≤ 100,000)에 대해 N - 1개의 연결 정보가 주어진다.
[3] 부모 노드 번호가 2번인 노드부터 순서대로 출력한다.

풀이

입력 예시가 아래와 같은 경우를 살펴보자.

노드와 그 연결 정보를 그래프(트리)로 그리면 다음과 같다.

여기서 1번 노드부터 시작하여 자식 노드를 찾고, 그것의 자식 노드를 또 찾아가는 식으로 접근을 했다.

그래서 DFS로 접근을 했고, 각 노드의 방문 여부를 저장하는 visited를 0으로 초기화한 뒤, 방문할 때 부모 노드의 값으로 업데이트했다.

예를 들어, 1번 노드에서 4번 노드를 방문했다면 visited[4]를 1로 업데이트하고, 이때 4번 노드를 부모 노드로 자식 노드를 찾는 과정을 반복하여 visited[2]와 visited[7]을 4로 업데이트했다.

visited = [0] * (n + 1)
def dfs(parent_node):
    for child_node in graph[parent_node]:
        if visited[child_node] == 0:
            visited[child_node] = parent_node

            dfs(child_node)

코드

import sys

# 런타임 에러(RecursionError) 방지 목적
sys.setrecursionlimit(10**6)


# 입력
n = int(sys.stdin.readline())

graph = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    a, b = map(int, sys.stdin.readline().split())

    graph[a].append(b)
    graph[b].append(a)

# print(graph)

# DFS
visited = [0] * (n + 1)
def dfs(parent_node):
    for child_node in graph[parent_node]:
        if visited[child_node] == 0:
            visited[child_node] = parent_node

            dfs(child_node)

# 탐색(부모 노드 찾기)
dfs(1)

# 출력
for node in visited[2:]:
    print(node)