[백준알고리즘/11725번/파이썬] 트리의 부모 찾기

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

11725번: 트리의 부모 찾기

 

문제: 루트가 없는 트리가 주어졌을 때, 루트를 1로 하여 각 노드의 부모 노드를 출력하는 문제 입니다.

 

풀이 방법: bfs 

 

풀이 코드:

import sys
import heapq

n = int(input())

graph = [[] for _ in range(n+1)]

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

data = []
visited = [False] * (n+1)

def bfs():
    q = [1]
    visited[1] = True
    while q:
        parent = heapq.heappop(q)
        for node in graph[parent]:
            if not visited[node]:
                visited[node] = True
                heapq.heappush(data,(node,parent))
                heapq.heappush(q,node)

bfs()

while data:
    print(heapq.heappop(data)[1])

 

풀이 설명:

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

먼저 방향성이 없는 노드의 2개가 주어지므로, 다음과 같이 graph에 양방향으로 저장해주었습니다.

 

data = []
visited = [False] * (n+1)

 

최종적으로 값을 출력하기 위한 data를 담기 위해서 data라는 Array를 선언해주었습니다.

그리고 bfs 탐색에서 방문 여부를 저장하기 위한 visited Array를 선언했습니다.

 

 bfs 동작하는 방식을 보면 다음과 같습니다.

def bfs():
    q = [1]
    visited[1] = True
    while q:
        parent = heapq.heappop(q)
        for node in graph[parent]:
            if not visited[node]:
                visited[node] = True
                heapq.heappush(data,(node,parent))
                heapq.heappush(q,node)

 

bfs를 통해 노드를 탐색했는데, 제일 핵심되는 부분은 탐색을 하면서 (현재 노드, 현재 노드의 부모 노드) 정보를 data라는 배열에 heapq 모듈을 통해 담은 것입니다.

 

최종 출력 값의 요구하는 점이 2번째 노드부터 순서대로 출력하는 것입니다.

 

따라서, heapq 모듈의 heapq.heappush()는 (현재 노드, 현재 노드의 부모 노드) 의 형식으로 저장하게 되면 현재 노드 기준 오름차순 정렬을 해서 저장을 해주므로 heapq 모듈을 활용하였습니다. (또한 log의 시간복잡도를 가지므로, 시간복잡도 측면에서도 우수하다는 이점도 있습니다.)

 

저희는 1을 루트로 하기로 했으므로, 큐의 역할을 하는 배열에 1을 넣고 탐색을 시작합니다.

 

for node in graph[parent]:
    if not visited[node]:
        visited[node] = True
        heapq.heappush(data,(node,parent))
        heapq.heappush(q,node)

1부터 노드를 탐색하며 방문하지 않은 노드에 대해서 처리를 진행해주는데 graph의 index로 들어가는 값이 부모 노드, graph[index] 안에 담긴 값이 index의 자식 노드 라고 볼 수 있습니다.

 

이를 바탕으로 앞서 말씀드린 것과 같이, 이와 같이 탐색과 데이터 저장을 동시에 진행하였습니다.

 

while data:
    print(heapq.heappop(data)[1])

마지막으로 다음과 같이 최종 값들을 출력해주었습니다.