[백준 1167번] 트리의 지름 (Python3)

import sys
input = sys.stdin.readline
from collections import deque
import heapq
INF = sys.maxsize

N = int(input())

graph = [[] for i in range(N+1)]

for i in range(N):
  data = list(map(int,input().split()))
  for j in range(1,len(data)-1):
    if j%2 == 0:
      continue
    graph[data[0]].append((data[j],data[j+1]))

hq = []
distance = [[-INF]*(N+1) for i in range(2)]

def Dijkstra(node,distance):
  heapq.heappush(hq,(0,node))

  while hq:
    w,now = heapq.heappop(hq)
    if distance[now] >= w:
      continue
    distance[now] = w
    
    for i in graph[now]:
      x,weight = i
      heapq.heappush(hq,(w-weight,x))

Dijkstra(1,distance[0])
Dijkstra(1+distance[0][1:].index(min(distance[0][1:])),distance[1])
  
print(-min(distance[1][1:]))

 

 

지난번에 풀었던 1967번 트리의 지름과 입력방식만 다르고 사실상 똑같은 문제다.

그래서 똑같은 알고리즘(다 익스트라)에서 입력방식만 바꿔줬는데, 시간 초과가 떴다...

뭐가 문제지 하다가 "트리기 때문에 굳이 다익스트라를 할 필요가 없다" 라는 조언으로 코드를 살짝 바꿔주었다.

 

import sys
input = sys.stdin.readline
from collections import deque
import heapq
INF = 10**9

N = int(input())

graph = [[] for i in range(N+1)]

for i in range(N):
  data = list(map(int,input().split()))
  for j in range(1,len(data)-1):
    if j%2 == 0:
      continue
    graph[data[0]].append((data[j],data[j+1]))

def BFS(start):
  Dist = [0]*(N+1)
  dq = deque()
  dq.append((start,0))

  while dq:
    now,w = dq.popleft()
    if Dist[now] == 0:
      Dist[now] = w
  
      for x1,w1 in graph[now]:
        if x1 != start:
          dq.append((x1,w1+w))
  return max(Dist),Dist.index(max(Dist))

l,x = BFS(1)

result, x = BFS(x)

print(result)

다 익스트라랑 거의 똑같은데 heapq를 쓰냐 dq를 쓰냐 그차이다. 트리니까 굳이 heapq를 쓸 필요가 없기 때문에 (어차피 경로가 하나밖에 없으므로)

 

그런데 문제는 이 코드는 틀렸습니다 가 떴었다. 그러다 질문게시판을 찾던 도중 알게된 사실은

 

"정점 번호는 순서대로 주어지지 않는다" 였다!

 

아니 어차피 트리라서 모든정점을 데이터로 줄건데 왜 이딴식으로 꼰건지 출제자의 의도가 이해가 안가지만 어쨌든 저걸 수정하니 정답이 맞았다.

그래서 혹시? 싶어서 다 익스트라 알고리즘도 시도해보니 이것도 시간만 아주 살짝 더 걸릴뿐 잘 통과되더라

왜 틀렸습니다가 아니라 시간초과가 뜬거야 근데...

 

 

결론은 위의 코드와 아래 코드 둘 다 잘 작동한다.

 

 

오늘의 교훈) 데이터 입력 방식을 꼼꼼이 검토하자