K번째 최단 거리 찾기 파이썬, 백준 1854 파이썬

K번째 최단 거리 찾기

개요

다익스트라 알고리즘을 이용하여 K번째 최단 거리를 구합니다.

구현 방법

• 최단 경로 리스트를 2차원 리스트로 설정하여 내부를 K개의 최단 경로로 채웁니다.
• K번째 경로를 찾기 위해서 여러 번 노드가 쓰일 수도 있으므로 방문 리스트는 따로 적용하지 않습니다.
• 위의 내용을 기억하며 아래의 최단 거리 리스트 채우기 규칙을 따릅니다.

1. 최단 경로 리스트 및 그래프 인접 리스트를 초기화합니다.
2. 우선 순위 큐에서 연결된 노드와 가중치 데이터를 가져옵니다.
3. 연결 노드의 K번째(마지막) 경로와 신규 경로를 비교해 신규 경로가 더 작을 때 K번째 경로에 값을 업데이트합니다. 이때 경로가 업데이트되는 경우 거리 배열을 오름차순으로 정렬하고 우선순위 큐에 연결 노드를 추가합니다.
4. 큐가 비워질때까지 2~3을 반복합니다.

백준 1854

import sys
input = sys.stdin.readline
print = sys.stdout.write

from queue import PriorityQueue

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

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

# 최단 거리 리스트를 K 번째 경로 개수만큼 2차원 배열로 설정
sp_list = [[3*(10**9) for j in range(K)] for i in range(N+1)]

for i in range(M):
    u, v, w = map(int, input().split())
    n_list[u].append([v, w])

my_pq = PriorityQueue()

sp_list[1][0] = 0
my_pq.put([0, 1])

while not my_pq.empty():
    cur_node_set = my_pq.get()
    sp = cur_node_set[0]
    cur_node = cur_node_set[1]
    for new_node_set in n_list[cur_node]:
        new_node = new_node_set[0]
        weight = new_node_set[1]
        # 마지막 최단 경로 값보다 새롭게 발견된 경로 값이 작으면 대체한 후 sort
        if sp_list[new_node][-1] > weight + sp:
            sp_list[new_node][-1] = weight+sp
            sp_list[new_node].sort()
            my_pq.put([weight+sp, new_node])

for i in range(1, N+1):
    if sp_list[i][K-1] == 3*(10**9):
        print('-1\n')
    else:
        print(str(sp_list[i][K-1]) + '\n')