[알고리즘] BOJ 1167 트리의 지름 #Python

[BOJ] 1167 트리의 지름 바로가기

📍 문제

트리의 지름이란, 트리에서 임의의 두 점 사이의 거리 중 가장 긴 것을 말한다. 트리의 지름을 구하는 프로그램을 작성하시오.

---

📍 입력

트리가 입력으로 주어진다. 먼저 첫 번째 줄에서는 트리의 정점의 개수 V가 주어지고 (2 ≤ V ≤ 100,000)둘째 줄부터 V개의 줄에 걸쳐 간선의 정보가 다음과 같이 주어진다. 정점 번호는 1부터 V까지 매겨져 있다.

먼저 정점 번호가 주어지고, 이어서 연결된 간선의 정보를 의미하는 정수가 두 개씩 주어지는데, 하나는 정점번호, 다른 하나는 그 정점까지의 거리이다. 예를 들어 네 번째 줄의 경우 정점 3은 정점 1과 거리가 2인 간선으로 연결되어 있고, 정점 4와는 거리가 3인 간선으로 연결되어 있는 것을 보여준다. 각 줄의 마지막에는 -1이 입력으로 주어진다. 주어지는 거리는 모두 10,000 이하의 자연수이다.

---

📍 출력

첫째 줄에 트리의 지름을 출력한다.

---

📍 풀이

🧷 풀이 과정

DFS 탐색으로 전수 조사하는 방식으로 코드를 작성한 결과 시간초과 오류가 발생했다.
방법을 알 수 없어서 해결 방식을 참고하기 위해 구글링을 한 결과 트리의 지름을 구하는 정형화된 방식이 있다는 것을 알게 되었다.

트리의 지름을 구하는 방식은 다음과 같다.

1. 임의의 정점 x에서 가장 먼 정점 y를 찾는다.
2. 정점 y에서 가장 먼 정점 z를 찾는다.
3. 정점 y에서 z까지의 거리가 트리의 지름이 된다.

트리의 지름을 구하는 방식에 대한 원리는 📒[트리의 지름 구하기] 에서 확인할 수 있다.

✍ 코드

# BOJ 1167 트리의 지름
# https://www.acmicpc.net/problem/1167

from sys import stdin
from collections import defaultdict, deque

def solution(V, edges):

    # 그래프 초기화
    graph = defaultdict(list)
    for edge in edges:
        for i in range(1, len(edge)-2, 2):
            graph[edge[0]].append(((edge[i], edge[i+1])))

    # BFS 탐색
    def bfs(startNode):
        visited = [-1] * (V + 1)
        visited[startNode] = 0
        queue = deque([startNode])
        while queue:
            node = queue.popleft()
            for nextNode, weight in graph[node]:
                if visited[nextNode] == -1:
                    visited[nextNode] = visited[node] + weight
                    queue.append(nextNode)

        maxWeight = max(visited)
        endNode = visited.index(maxWeight)
        return maxWeight, endNode
    
    # 1. 임의의 정점 x에서 가장 먼 정점 y를 찾는다.
    maxWeight, y = bfs(1)
    # 2. 정점 y에서 가장 먼 정점 z를 찾는다.
    maxWeight, z = bfs(y)

    return maxWeight

# input
V = int(stdin.readline())
edges = list(list(map(int,stdin.readline().split())) for _ in range(V))

# response
response = solution(V, edges)
print(response)