[백준/BOJ] platinum4 - 1854번 K번째 최단경로 찾기 (Python)

▶1854 - K번째 최단경로 찾기

▶문제

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

 

▶입력

첫째 줄에 n, m, k가 주어진다. (1 ≤ n ≤ 1000, 0 ≤ m ≤ 2000000, 1 ≤ k ≤ 100) n과 m은 각각 김 조교가 여행을 고려하고 있는 도시들의 개수와, 도시 간에 존재하는 도로의 수이다.

이어지는 m개의 줄에는 각각 도로의 정보를 제공하는 세 개의 정수 a, b, c가 포함되어 있다. 이것은 a번 도시에서 b번 도시로 갈 때는 c의 시간이 걸린다는 의미이다. (1 ≤ a, b ≤ n. 1 ≤ c ≤ 1000)

도시의 번호는 1번부터 n번까지 연속하여 붙어 있으며, 1번 도시는 시작 도시이다.

 

▶출력

n개의 줄을 출력한다. i번째 줄에 1번 도시에서 i번 도시로 가는 k번째 최단경로의 소요시간을 출력한다.

경로의 소요시간은 경로 위에 있는 도로들을 따라 이동하는데 필요한 시간들의 합이다. i번 도시에서 i번 도시로 가는 최단경로는 0이지만, 일반적인 k번째 최단경로는 0이 아닐 수 있음에 유의한다. 또, k번째 최단경로가 존재하지 않으면 -1을 출력한다. 최단경로에 같은 정점이 여러 번 포함되어도 된다.

 

▶풀이

단순한 다익스트라 문제이다. 그냥 평소에 풀던 거처럼 풀면 되는데, 한 가지 추가되는 사항이 있다.

k번째 최단 경로이다. 그래서 원래 dp는 1차원이었다면, 여기서는 k차원으로 만들어 주었다.

 

그럼 여기서 어떻게 넣어줄 것인가?

k-1(마지막) 번째 값을 비교해주면 된다.

 

k가 2라고 가정하면 i번째 노드엔 초기값 99999 99999가 들어있다.

여기서 처음에 i까지 가는 길이가 10이 나오면 99999 10 이렇게 바뀐다.

그다음 10 99999로 정렬해주고, 다음에 2가 들어오면 2 10 이렇게 될 것이다.

 

어떤 수가 들어와도 가장 작은 값이 0번째, 두 번째 작은 값이 1번째에 들어오게 될 것이다.

k가 커지더라도 똑같은 방식으로 들어오게 된다.

 

그다음 각각 노드들의 k-1번째 값을 출력해주면 된다.

import sys
import heapq


def dijkstra(start):
    dp[start][0] = 0
    queue = []
    heapq.heappush(queue, [0, start])
    while queue:
        cur_dist, node = heapq.heappop(queue)
        for ad_node, d in graph[node]:
            w_dist = cur_dist + d
            if w_dist < dp[ad_node][k - 1]:
                dp[ad_node][k-1] = w_dist
                dp[ad_node].sort()
                heapq.heappush(queue, [w_dist, ad_node])


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

for i in range(1, n + 1):
    if dp[i][k-1] == sys.maxsize:
        print(-1)
    else:
        print(dp[i][k-1])

한 가지 흠이 있다면, 이번 문제 또한 python3으로 제출해서 성공하지 못했다.

pypy3로는 성공했지만, python3으로는 시간 초과가 떴다.

아무래도 python3으로 다익스트라 공부 좀 해야 할 것 같다.