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

본문 링크

import sys,heapq
input=sys.stdin.readline
INF=int(1e9)

def Dijkstra():

    while heap:
        value,node=heapq.heappop(heap)

        for next_value,next_node in graph[node]:

            next_value+=value

            if len(dp[next_node])<K:
                heapq.heappush(dp[next_node] , -next_value)
                heapq.heappush(heap, (next_value , next_node))
            elif next_value<-dp[next_node][0]:
                # dp 해당점점에서 최대힙의 인덱스 0번째값, 즉 최대값이 next_value 보다 크다면
                # 이동을 한다. 이때 dp의 해당정점을 방문한 개수는 K개 이상이다.
                # 즉 해당정점을 방문한 개수가 K개 이상이고 그때의 최대값이 이번에 이동할 값보다 크다면
                # 그 최대값을 삭제하고 더 작은 값을 넣어준다.
                # 따라서 pop 을 먼저해주고 push 를 한다.
                heapq.heappop(dp[next_node])
                heapq.heappush(dp[next_node] , -next_value)
                heapq.heappush(heap, (next_value , next_node))


N,M,K=map(int,input().split())

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

for i in range(M):
    a,b,c=map(int,input().split())
    graph[a].append((c,b)) #단방향 , 가중치먼저

# 전처리
dp=[ [] for _ in range(N+1) ]
dp[1].append(0)
heap=[] ; heapq.heappush(heap,(0,1))

Dijkstra() #함수 실행
for i in range(1,N+1):
    if len(dp[i])<K:
        print(-1)
    else:
        print(-dp[i][0])

📌 어떻게 접근할 것인가?

다익스트라 알고리즘을 사용했습니다. 문제에서 구하고자 하는 것은 $K$ 번째 최단경로 입니다.

기존에 사용하던 dp 배열 , 즉 최단경로를 담는 배열을 2차원 배열로 만들어 주고
최대 힙으로 구성하였습니다.

이렇게 한 이유는 그냥 $K$번째 최단 경로를 일반적인 다익스트라로 구현했을때
꼭 $K$ 번째 최단경로가 모든 경로를 구한 후 정렬했을때 $K$ 번째가 되지 않을 수 있기 때문입니다.

실제로 최단경로를 구하는 과정중에 첫번째로 들어오는 수가 두번째로 들어오는 수 보다 클 수 있습니다.

따라서 각 인덱스(노드의 번호) 마다 이동했을때의 경로의 합을 넣어주고 만약 그 개수가 $K$ 개 이상이면 최대힙을 사용하여 최대값을 빼주고 더 적은 값을 넣어줍니다.

따라서 모든 노드에 대해서 방문하면 최대힙의 인덱스 0번째 값이 모든 정점들에 대해 방문했을때 $K$ 번쨰 값이 됩니다.