알고리즘 설계 실습 - Dijkstra algorithm

---

문제

상원이는 네비게이션 개발자입니다. 상원이는 사용자들이 다양한 경로로 가기를 원해서 가장 빠르게 갈 수 있는 최단 경로가 아니라, k번째로 빠르게 갈 수 있는 경로를 구해서 사용자들에게 제  하려합니다. n개의 위치와 m개의 도로에 대한 정보가 주어졌을 때, k번째로 빠르게 도착할 수 있는 경로를 구하는 프로그램을 작성하세요.

입력

첫째 줄에 n, m, k가 주어집니다. (1≤n≤1,000, 0≤m≤250,000, 2≤k≤50) n과 m은 각각 위치들의 개수와, 위치 간에 존재하는 도로의 수입니다. 이어지는 m개의 줄에는 각각 도로의 정보를 제공하는 세 개의 정수 a, b, c가 주어집니다. 이것은 a번 위치에서 b번 위치로 이동하는 도로를 사용할 때, c의 시간이 걸린다는 의미입니다. (1≤a, b≤n, 1≤c≤1,000) a와 b는 위치의 번호이며, 시작 도시는 1번 위치입니다. 두 도로의 시작점과 도착점이 모두 같은 경우는 없습니다.

출력

n개의 줄을 출력해야 합니다. i번째 줄에 1번 위치에서 i번 위치로 가는 k번째로 빠른 경로의 소요시간을 출력하세요. 경로의 소요시간은 경로 위에 있는 도로들을 따라 이동하는데 필요한 시간들의 합입니다. k번째 최단경로가 존재하지 않는 경우 -1을 출력하세요. 최단 경로에 목적지를 포함한 같은 위치가 여러 번 포함되어도 됩니다. 만약 경로는 다르지만 소요시간은 같은 경우 다른 경로로 취급합니다 (예시 2, 3, 4의 4번 위치로 가는 결과 참고). 1번 위치에서 1번 위치로 가는 1번째로 빠른 경로의 소요시간은 0입니다.

---

풀이

dijkstra.py

import sys
import heapq

input = sys.stdin.read().split()
INF = int(1e8)

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

# 각 노드별로 경로의 비용을 저장할 리스트
distance = [[] for _ in range(N+1)]
# 그래프의 간선 정보를 저장할 리스트
graph = [[] for _ in range(N+1)]

# 모든 간선 정보 입력 받기
for _ in range(M):
    a, b, c = map(int, input().split())
    graph[a].append((c, b))

# 시작 노드 1에 대해 비용 0을 우선순위 큐에 푸시
heapq.heappush(distance[1], 0)
# 우선순위 큐에 시작 노드와 비용을 푸시
heapq.heappush(queue, (0, 1))

# 우선순위 큐가 빌 때까지 반복
while queue:
    # 현재 노드와 비용을 팝
    dist, pos = heapq.heappop(queue)
    # 현재 노드와 인접한 노드들 검사
    for i in graph[pos]:
        cost = dist + i[0]
        # 해당 노드에 저장된 거리가 K개보다 적으면 비용을 푸시
        if len(distance[i[1]]) < K:
            heapq.heappush(distance[i[1]], -cost)
            heapq.heappush(queue, (cost, i[1]))
        # 이미 K개가 저장되어 있고, 새로운 비용이 K번째 비용보다 작으면 갱신
        elif cost < -distance[i[1]][0]:
            heapq.heappop(distance[i[1]])
            heapq.heappush(distance[i[1]], -cost)
            heapq.heappush(queue, (cost, i[1]))

# 모든 노드에 대해 K번째 최단 경로가 있는지 검사하고 출력
for i in range(1, N+1):
    if len(distance[i]) == K:
        ans = -distance[i][0]
        print(ans)
    else:
        print(-1)

이 코드는 주어진 그래프에서 각 노드로의 K번째 최단 경로를 찾기 위해 다익스트라 알고리즘을 변형하여 각 노드에 도달하는 모든 가능한 경로 중에서 비용이 K번째로 작은 경로를 찾는다. 그래프의 각 노드와 간선 정보는 리스트로 관리되며, 우선순위 큐를 사용하여 가장 짧은 경로를 우선적으로 처리한다. 경로의 비용은 heap를 이용하여 각 노드에 대해 K개의 최소 비용만을 유지한다. 마지막으로, 각 노드별로 K번째 최단 경로의 존재 여부에 따라 해당 비용을 출력하거나 경로가 존재하지 않는 경우 -1을 출력한다.

---