K번째 최단경로 찾기

K번째 최단경로 찾기

접근방법

heapq와 Dijkstra 알고리즘을 이용하였다

heapq(우선순위 큐)에 대하여 알아보자면,

모든 부모 노드가 자식보다 작거나 같은 값을 갖는 이진 트리이다.

파이썬 에서는 heappush 를 사용하여 큐에 push 할 수 있고,

heappop 을 사용하여 루트노드를 pop 할 수 있다.

다음과 같이 리스트로 넣는다면 첫번째 인덱스가 비교된다.

import heapq

q = []

heapq.heappush(q, [1, 99])
heapq.heappush(q, [3, 1])
heapq.heappush(q, [2, 55])

print(q)

print(heapq.heappop(q))

print(q)

[[1, 99], [3, 1], [2, 55]]
[1, 99]
[[2, 55], [3, 1]]

하지만 우리는 최소 거리부터 K 번째를 구해야 하기 때문에, 다음과 같이 첫번째 인덱스(총 거리) 를 음수로 넣어줄 것이다.

import heapq

qq = []

heapq.heappush(qq, [-1, 99])
heapq.heappush(qq, [-3, 1])
heapq.heappush(qq, [-2, 55])

print(qq)

[[-3, 1], [-1, 99], [-2, 55]]

위의 방식대로 Dijkstra 알고리즘에 적용하여 코드를 풀어나가보자면,

해당 노드까지 가는 거리의 갯수가 k개보다 작다면

리스트에 바로 삽입을 해주고

k개 이상이라면

루트노드(가장 먼 거리) 를 팝 해준후에 삽입을 해주면 총 경로 개수가 k개로 유지가 된다

이후 결과출력때 거리가 k개 라면 바로 프린트 해주고, k개가 되지 않는다면 -1을 프린트 해주면 된다.

코드

# 1854 / K변째 최단경로 찾기
# https://www.acmicpc.net/problem/1854
import sys
import heapq
input = sys.stdin.readline

n, m, k = map(int, input().split())
arr =[[] for _ in range(n + 1)]
for _ in range(m):
    a, b, c = map(int, input().split())
    arr[a].append([b, c])

result = [[] for _ in range(n + 1)]
result[1].append(0)
q = [[0, 1]]

while q:
    cost, node = heapq.heappop(q)

    for next, w in arr[node]:
        dis = cost + w
        if len(result[next]) < k:
            heapq.heappush(result[next], -dis)
            heapq.heappush(q, [dis, next])
        elif dis < -result[next][0]:
            heapq.heappop(result[next])
            heapq.heappush(result[next], -dis)
            heapq.heappush(q, [dis, next])

for idx in range(1, n + 1):
    if len(result[idx]) < k:
        print(-1)
    else:
        print(-result[idx][0])

#