[백준1967&파이썬] 트리의 지름을 다익스트라를 이용해 구하자

# 문제

백준 1967 트리의 지름 파이썬 풀이

1967번: 트리의 지름

# 코드

import sys
from collections import defaultdict
from heapq import *


def print2D(arr) :
  for i in arr :
    print(*i)

# 노드의 개수
n = int(sys.stdin.readline().strip())
trees = defaultdict(list)

distance = [0] * (n+1)
visit = [-1] * (n+1)
for _ in range(n-1) :
  a,b,c = map(int,sys.stdin.readline().split())
  # 양방향 연결
  trees[a].append((b,c))
  trees[b].append((a,c))

# 다익스트라로 1번노드에서 가장 먼 정점을 구한다.
# 그정점에서 다시 다익스트라를 사용해서 가장 거리가 먼곳을 구한다
# 근데 이건 BFS써도 될거 같은데 난 다익스트라 까먹어서 다익스트라로 풀어보겠다.

def dijkstra(start) :
  global n
  h = list()
  distance = [float('inf')] * (n+1)
  distance[start] = 0 # 시작점에서 시작점으로 출발하는 경우는 0
  heappush(h,(0,start))

  while len(h) != 0 :
    cur_cost, cur_node = heappop(h) # ( cost, node )
    for next_node, next_cost in trees[cur_node] :
      updated_cost = cur_cost + next_cost
      if updated_cost < distance[next_node] :
        distance[next_node] = updated_cost
        heappush(h,(updated_cost,next_node))

  return distance

start_distance = dijkstra(1)
print(max(dijkstra(start_distance.index(max(start_distance[1:])))[1:]))

# 풀이

다익스트라 안쓴지 오래되서 다익스트라로 풀었다. 근데 BFS나 DFS로 풀어도 상관 없다.

일단 시작노드(1)에서 가장 멀리있는 노드를 구한다.

가장 멀리 있는 노드가 구해지면 그 노드는 트리의 지름의 한 끝점이 된다.

그러면 한 끝점을 찾았으니까 그점에서 다익스트라 돌려서 가장 먼 곳을 구하여 지름을 구한다.

결국 이 문제의 핵심은 트리의 leaf를 구할 수 있으면 된다. 위 풀이처럼이 아니더라도 leaf노드 구한다음에 leaf노드로 다익스트라 돌려서 최대값 찾으면 될거 같다. ( 시간초과 날 수도 있을거 같긴 하다)

결국 리프노드에서 리프노드까지 가는게 무조건 최장거리이기 때문이다.