백준 1854 : K번째 최단경로 찾기 (Python)

출처 : https://www.acmicpc.net/problem/1854

import heapq

n, m, k = map(int, input().split())
graph = [[] for _ in range(n+1)]
for _ in range(m) :
    a, b, c = map(int, input().split())
    graph[a].append((c, b))   

dist = [[] for _ in range(n+1)]
def dijkstra(start, k) :
    # pq, dist에 초기값 삽입
    pq = []
    heapq.heappush(dist[start], 0)
    heapq.heappush(pq, (0, start))
    while pq :
        # dist_now : 시작점부터 현 노드까지의 거리
        dist_now, now = heapq.heappop(pq)
        # dist_now_near : 현 노드에서 인접한 이웃 노드까지의 거리
        for dist_now_near, near in graph[now] :
            # dist_near : 시작점부터 인접 노드까지의 거리
            dist_near = dist_now + dist_now_near
            # dist[near]에 저장된 거리가 k개 미만이면 그대로 저장
            if len(dist[near]) < k :
                heapq.heappush(dist[near], -dist_near)
                heapq.heappush(pq, (dist_near, near))
            else :
                # 저장된 거리 중 가장 큰 값 대체
                if dist_near < -dist[near][0] :
                    heapq.heappop(dist[near])
                    heapq.heappush(dist[near], -dist_near)
                    heapq.heappush(pq, (dist_near, near))

dijkstra(1, k)
for i in range(1, n+1) :
    if len(dist[i]) < k :
        print(-1)
    else :
        print(-dist[i][0])

• k번째의 최단거리를 찾는 문제
• pq : 일반적인 다익스트라처럼 갱신된 (거리, 노드)를 저장하는 우선순위큐(최소)
• dist[node] : node에 도달할 수 있는 모든 경로의 거리를 최대 k개 저장한 우선순위큐(최대)
  • 가능한 모든 거리 중 가장 작은 k개를 저장
• 현재 dist[node]에 저장된 거리가 k개 미만이면 dist_near 그대로 저장
• 이미 dist[node]에 k개의 거리가 저장된 경우
  • 저장된 거리 중 가장 큰 거리(-dist[near][0])가 dist_near 보다 크다면 대체
• 완성된 dist 배열에서 길이가 k미만인 dist[node]는 node까지 도달하는 k번째 최단거리가 없다는 뜻이므로 -1 출력