[BFS] 백준 #1167 트리의 지름 구하기 - Python

📚 문제

링크: https://www.acmicpc.net/problem/1167

1167번: 트리의 지름

 

💡 접근법

BFS를 이용하여 각 노드까지의 거리를 구하고, 이 중 최대 거리를 갖는 노드에서 다시 BFS를 수행한다.

이때 각 노드까지의 거리 중 최대 거리가 바로 트리의 지름이다.

 

즉, 이 문제의 아이디어는 "임의의 노드에서 가장 긴 경로로 연결돼 있는 노드는 트리의 지름에 해당하는 두 노드 중 하나다." 이다.

BFS란?
시작 노드에서 출발해 시작 노드를 기준으로 가까운 노드를 먼저 방문하면서 탐색하는 알고리즘이다.

 

💻 코드

1) 전체 코드

import sys
from collections import deque

N = int(sys.stdin.readline().strip())
A = [[] for _ in range(N+1)]
visited = [False] * (N+1)
distance = [0] * (N+1)

for _ in range(N):
    temp = list(map(int, sys.stdin.readline().split()))
    node = temp[0]
    data = temp[1:]
    for i in range(0, len(data), 2):
        if data[i] == -1:
            break
        A[node].append((data[i], data[i+1]))


def BFS(v):
    queue = deque()
    queue.append(v)
    visited[v] = True
    while queue:
        node = queue.popleft()
        for i in A[node]:
            if not visited[i[0]]:
                visited[i[0]] = True
                distance[i[0]] = distance[node] + i[1]
                queue.append(i[0])


BFS(1)
max_idx = 1
for i in range(2, N+1):
    if distance[max_idx] < distance[i]:
        max_idx = i

visited = [False] * (N+1)
distance = [0] * (N+1)
BFS(max_idx)

print(max(distance))

 

2) 해설

(1) 패키지 Import

import sys
from collections import deque

파이썬 내장 함수인 input()보다 빠르게 값을 입력받기 위해 sys 라이브러리, 큐를 사용하기 위해 collections.deque 모듈을 임포트한다.

 

(2) 입력받기

N = int(sys.stdin.readline().strip())
A = [[] for _ in range(N+1)]
visited = [False] * (N+1)
distance = [0] * (N+1)

for _ in range(N):
    temp = list(map(int, sys.stdin.readline().split()))
    node = temp[0]
    data = temp[1:]
    for i in range(0, len(data), 2):
        if data[i] == -1:
            break
        A[node].append((data[i], data[i+1]))

트리의 노드 개수(N)를 입력받고, 인접 리스트(A)와 방문 리스트(visited)를 초기화한다.

해당 문제에서는 특정 노드에서부터 각 노드까지의 거리를 저장하는 distance 리스트가 추가로 필요하다. 이 또한 각 요소를 0으로 초기화한다.

 

노드의 개수만큼 for 문을 순회하면서 인접 리스트를 생성한다.

 

(3) BFS 구현

def BFS(v):
    queue = deque()
    queue.append(v)
    visited[v] = True
    while queue:
        node = queue.popleft()
        for i in A[node]:
            if not visited[i[0]]:
            	queue.append(i[0])
                visited[i[0]] = True
                distance[i[0]] = distance[node] + i[1]

deque()를 생성하고 현재 노드를 큐에 삽입한다.

현재 노드에 대해 방문 리스트를 체크하고 큐가 빌 때까지 while 문을 수행하기 시작한다.

 

큐에서 노드를 꺼내고 해당 노드의 인접 리스트의 요소들 중 방문하지 않은 노드를 큐에 삽입하고 방문 리스트를 체크한다.

그리고 distance[i[0]]에 특정 노드로부터 직전 노드까지의 거리 + 현재 노드의 가중치를 저장한다.

 

(4) 출력하기

BFS(1)
max_idx = 1
for i in range(2, N+1):
    if distance[max_idx] < distance[i]:
        max_idx = i

visited = [False] * (N+1)
distance = [0] * (N+1)
BFS(max_idx)

print(max(distance))

노드 1을 시작점으로 하여 BFS를 수행한다.

BFS 수행 결과, distance 리스트에 각 노드까지의 거리가 저장되었을 것이다. 최대 거리가 저장된 노드를 찾아낸다.

해당 노드가 트리에 지름에 해당하는 두 노드 중 하나이다.

 

visited와 distance 리스트를 초기화해주고 최대 거리가 저장됐던 노드(트리의 지름에 해당하는 두 노드 중 하나)를 시작점으로 하여 다시 BFS를 수행한다.

BFS 수행 이후, distance 리스트에서 최댓값을 출력한다. 해당 값이 트리의 지름이다.

 

✅ 정답 확인

 

---

 

포스팅한 내용에 대한 조언/지적/피드백 모두 환영입니다!

읽어주셔서 감사합니다^.^