[백준 1854번] K번째 최단경로 찾기 (Python3)

import sys
input = sys.stdin.readline
from heapq import *

N,M,K = map(int,input().split())
graph = [[] for i in range(N+1)]

for _ in range(M):
  x,y,w = map(int,input().split())
  graph[x].append((y,w))

hq = [(0,1)]; distance = [[] for i in range(N+1)]
while hq:
  w,now = heappop(hq)
  if len(distance[now])==K:
    if w >= -distance[now][0]:
      continue
    else:
      heappop(distance[now])
  heappush(distance[now],-w)
  for next,w1 in graph[now]:
    heappush(hq,(w+w1,next))

for i in range(1,N+1):
  print(-1 if len(distance[i])<K else -distance[i][0])

다 익스트라 응용문제

아이디어를 떠올리기 꽤 어려웠다.

 

내가 처음에 생각한 방법은, 다 익스트라로 최단거리를 미리 측정한 뒤, 다시 다 익스트라를 하여 최단거리와 그렇지 않은 거리간의 차이를 구하고, 이를 DFS 역추적으로 구하는 방식이었다.

그러나 이 방법의 문제는 사이클이 존재할 수 있다는 것이었다. 사이클이 존재하면 DFS 과정에서 무한루프가 발생하는 것이었다.

 

따라서 어떻게 해결할까 고민하다가 채점현황을 눌러봤는데, 생각보다 통과한 코드들의 길이가 짧았다. 내가 너무 어렵게 생각했다는걸 깨닫고 좀더 심플한 아이디어가 뭘까 생각하다가, 그냥 다 익스트라 자체에서 K번째 경로를 찾아야 한다는 것을 알게되었다.

 

핵심 아이디어는, 다 익스트라에서 DP를 1차원 리스트가 아닌 2차원 최대힙으로 두는 것이다. 그리고 경로를 갱신할 때, DP[now]보다 작은 경우에 갱신하는 것이 아니라 DP[now]의 길이가 K보다 작거나, K번째 값보다 현재 값이 작은 경우에 갱신한다. 이러한 방법으로 다 익스트라 1번으로 K번째 최단거리를 구할 수 있다.

 

 

오늘의 교훈) 너무 복잡하게 생각하지 말자