[백준] 1854번: K번째 최단경로 찾기 with Python

BOJ 1854

• Dijkstra

문제

---

봄캠프를 마친 김진영 조교는 여러 도시를 돌며 여행을 다닐 계획이다. 그런데 김 조교는, '느림의 미학'을 중요시하는 사람이라 항상 최단경로로만 이동하는 것은 별로 좋아하지 않는다. 하지만 너무 시간이 오래 걸리는 경로도 그리 매력적인 것만은 아니어서, 적당한 타협안인 'k번째 최단경로'를 구하길 원한다. 그를 돕기 위한 프로그램을 작성해 보자.

풀이과정

---

일반적인 다익스트라 알고리즘과 유사하지만 K번 째 최단경로를 구해야한다는 점에서 차이가 있다. 기존, 다익스트라 알고리즘에서는 i까지의 최단거리를 distance를 저장하게 된다. 하지만 이 문제에서는 K번 째 최단경로를 구해야하므로, 가장 작은 값부터 k번째 최단경로까지를 모두 저장해야한다. 즉, distance는 i까지의 최단거리를 저장하는 2차원 배열이 되게 된다.

k번째 최단거리를 갱신할 때마다 i번째 도시까지 걸리는 가장 distance 리스트 내에서 가장 큰 값을 탐색하려고 하면 O(N) 시간이 걸리므로 시간복잡도가 올라가게 된다. 따라서, distance를 단순한 2차원 배열이 아닌 우선순위 큐(heapq)의 배열로 구현했다.

우선순위 큐의 길이는 K로 고정해야 한다. 따라서 다익스트라 알고리즘 대로 진행할 때 큐의 길이에 따라 코드의 실행 과정이 바뀐다.

		# 저장된 거리가 K보다 작을 때
        if len(distance[i[1]]) < K:
            # 그냥 넣기
            heapq.heappush(distance[i[1]], -cost)
            heapq.heappush(q, (cost, i[1]))
        # 저장된 거리가 K보다 크고
        # heapq의 가장 작은 값에 -1을 곱한 값보다 cost가 작을 때
        elif cost < -distance[i[1]][0]:
            # heapq에서 데이터 하나를 뺀다
            heapq.heappop(distance[i[1]])
            heapq.heappush(distance[i[1]], -cost)
            heapq.heappush(q, (cost, i[1]))

큐의 길이가 K보다 작다면 K번 째 작은 값이 존재하지 않으므로 cost 값에 관계없이 큐에 넣는다. 반면, 큐의 길이가 K보다 큰 경우에는 cost의 값이 큐 안에 존재하는 가장 큰 값보다 작을 때, 가장 큰 값을 제거하고 넣어줘야 한다. 파이썬의 heapq는 MIN Queue이므로 -1을 곱해서 큐에 값을 넣었다.

코드

---

import sys
import heapq

input = sys.stdin.readline
INF = int(1e12)

# 1 ≤ n ≤ 1000, 0 ≤ m ≤ 2000000, 1 ≤ k ≤ 100
N, M, K = map(int, input().split())

# 거리, 위치
q = []
# a번 도시에서 b번 도시로 갈 때는 c의 시간이 걸린다. (1 ≤ a, b ≤ n. 1 ≤ c ≤ 1000)

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))

heapq.heappush(distance[1], 0)  # 비용
heapq.heappush(q, (0, 1))       # 비용, 장소

while q:
    dist, now = heapq.heappop(q)
    for i in graph[now]:
        # i[0]: 비용, i[1]: 장소
        cost = dist + i[0]
        # 저장된 거리가 K보다 작을 때
        if len(distance[i[1]]) < K:
            # 그냥 넣기
            heapq.heappush(distance[i[1]], -cost)
            heapq.heappush(q, (cost, i[1]))
        # 저장된 거리가 K보다 크고
        # heapq의 가장 작은 값에 -1을 곱한 값보다 cost가 작을 때
        elif cost < -distance[i[1]][0]:
            # heapq에서 데이터 하나를 뺀다
            heapq.heappop(distance[i[1]])
            heapq.heappush(distance[i[1]], -cost)
            heapq.heappush(q, (cost, i[1]))

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