백준 1854 : K번째 최단경로 찾기

https://www.acmicpc.net/problem/1854

1854번: K번째 최단경로 찾기

 

해결 방법 : 다익스트라(Dijkstra)

• 다익스트라 알고리즘을 조금 응용하면 된다.
• 최단 경로를 저장하는 배열 fees에 등장 횟수를 나타내는 count 변수를 추가하여, count가 K가 될 때까지 fees의 value를 갱신하면 된다.
  • 다익스트라를 min-heap으로 구현하면 pop연산에 따라 현재 가장 작은 cost를 가진 경로가 나오기 때문에 특정 노드가 K번 pop 되었다면 당시의 cost가 K번째 최단 경로 비용이 될 것이다.
• 사이클이 존재할 경우 무한히 돌지 않도록 해당 노드가 K번 등장했다면 skip 해주자.
• 추가적인 설명으로는, heap push하는 쪽에서 k를 count해주는 것은 의미가 없다. (실질적으로 pop 연산으로 나올 때, 해당 값이 최단 경로인 것이기 때문이다.)
  • 하지만 heap push 전에 count == k라면 skip 해주는 코드는 시간을 절약시켜 줄 수 있다.

 

import sys
input = sys.stdin.readline
from collections import defaultdict
import heapq

def solution(n, k, G):
    fees = [[float('inf'), 0] for _ in range(n + 1)] # [count 번째 최소 비용, count]
    Q = [(0, 1)] # min heap
    
    while Q: # Dijkstra
        cur_f, cur_n = heapq.heappop(Q)
        if fees[cur_n][1] == k: continue # 만약 현재 노드가 이미 k번 등장했으면 skip
        fees[cur_n][0] = cur_f # 현재 node에서의 count 번째 최소 비용 갱신
        fees[cur_n][1] += 1 # 현재 node의 등장 count ++

        for next_n, next_f in G[cur_n]:
            if fees[next_n][1] == k: continue # 시간 복잡도 절약(안써도 통과는 된다.)
            heapq.heappush(Q, (cur_f + next_f, next_n))

    return [str(fee) if c == k else '-1' for fee, c in fees[1:]]

if __name__ == '__main__':
    n, m, k = map(int, input().split())
    G = defaultdict(list)
    for _ in range(m):
        f, t, c = map(int, input().split())
        G[f].append((t, c))
    result = solution(n, k, G)
    print('\n'.join(result))