[K번째 최단경로 찾기] 다익스트라 | 백준 1854번 | python

 

다익스트라 알고리즘에 대해 [이전글]에서 어떻게 풀어야하는지 정리해뒀다.

요약하자면!

[이전글] 에서는 인접 노드의 distance > 현재 노드의 distance + 인접 노드까지의 distance 를 비교해야 했다.

이를 위해서 우선순위 큐(PriorityQueue)와, 노드 방문 여부, disatnce를 고려해야 한다.

하지만 이번 문제에선 heapq를 사용한다.

파이썬 내장 heapq도 [이전글]에서 정리했다.

---

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

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

예제 입력1

5 10 2
1 2 2
1 3 7
1 4 5
1 5 6
2 4 2
2 3 4
3 4 6
3 5 8
5 2 4
5 4 1

예제 출력1

-1
10
7
5
14

 

📍 풀이

이전과는 다르게 k번째 거리를 구해야한다. 즉, 최단 거리를 구할 때는 visited로 이전 방문 노드를 기록해놨었지만, k번째 거리를 구하려면 방문 했던 곳을 또 방문을 해서 거리를 비교해야한다.
결론 visited를 기록할 필요가 없다!
또한 k개의 거리를 저장할 distance를 선언해야 한다.

그리고 우선순위 큐를 사용할 것이다.

우선순위 큐로 선언하면 편리한점

다익스트라 알고리즘 수행을 위한 노드 데이터를 저장하는 객체 형식을 우선순위 큐로 선언했기 때문에 새로운 노드가 삽입됐을 때 별도의 정렬을 해주지 않아도 자동으로 정렬돼 편리하게 구현할 수 있다는 장점이 있습니다.

힌트

초기 설정

• direction
• distance
• pq

 

direction
경로와 시간을 저장

direction = [[] for _ in range(n + 1)]
for _ in range(m):
    start, end, time = map(int, input().split())

    direction[start].append([end, time])
    
#결과
# 1  ->  [[2, 2], [3, 7], [4, 5], [5, 6]]
# 2  ->  [[4, 2], [3, 4]]
# 3  ->  [[4, 6], [5, 8]]
# 4  ->  []
# 5  ->  [[2, 4], [4, 1]]

distance
k 개의 row를 갖는 2차원 리스트 형태

distance = [[sys.maxsize] * k for _ in range(n + 1)]
distance[1][0] = 0

# 결과
# 1  ->  0 9223372036854775807 
# 2  ->  9223372036854775807 9223372036854775807 
# 3  ->  9223372036854775807 9223372036854775807 
# 4  ->  9223372036854775807 9223372036854775807 
# 5  ->  9223372036854775807 9223372036854775807 

pq
(현재 노드의 거리) + (다음 노드까지의 거리) < (다음 노드의 distance) 갱신

pq = [(0, 1)]

 

⚒️ 다익스트라 수행

heapq에 삽입할 때는 [거리, 노드] 순으로 넣어야한다.

while pq:
    cost, node = heapq.heappop(pq)
    
    for nextNode, nextCost in direction[node]:
        sumCost = cost + nextCost

        if distance[nextNode][k - 1] > sumCost:
            distance[nextNode][k - 1] = sumCost
            distance[nextNode].sort()

            heapq.heappush(pq, [sumCost, nextNode])

흐름

1️⃣ 1번 노드에서 갈 수 있는 노드 갱신

• 힙 큐 = [[0,1]]
• pq에서 pop한 노드 = 1, cost = 0
• 노드 1에서 갈 수 있는 노드, 현재 cost(0) + 인접노드까지의 거리 정보 삽입 ( [2,2], [5, 4], [6, 5], [7, 3] )
• 결과 : 힙 큐 =  [[2,2], [5, 4], [6, 5], [7, 3]] 가장 짧은 거리 위주로 정렬

distance 결과

node	[0]	[1]
1	0	∞
2	2	∞
3	7	∞
4	5	∞
5	6	∞

2️⃣ 2번 노드에서 갈 수 있는 노드 갱신 ⭐️

• 힙 큐 = [[2,2], [5, 4], [6, 5], [7, 3]]
• pq에서 pop한 노드 = 2, cost = 2
• 노드 2에서 갈 수 있는 노드, 현재 cost(2) + 인접노드까지의 거리 정보 삽입 ([6,3], [4,4])
• 결과 : 힙 큐 = [[4, 4], [5, 4], [6, 3], [7, 3], [6, 5]] 가장 짧은 거리 위주로 정렬

1. 노드 2에서 노드 4까지 가는 경우,
   ∞(distance[4][1]) > 2(현재 cost) + 2(노드2에서 노드4까지 가는 cost) 이므로
   distance[4][1] = 4으로 변경하고 distance[4]를 정렬한다.
   정렬하면 distance[4][0]인 5와 distance[4][1]인 4 값이 오름차순으로 정렬된다.
   
2. 노드 2에서 노드3까지 가는 경우,
   ∞(distance[3][1]) > 2(현재 cost) + 4(노드2에서 노드3까지 가는 nextCost) 이므로
   distance[3][1] = 6으로 변경하고 distance[3]을 정렬한다.
   정렬하면 distance[3][0]인 7와 distance[3][1]인 6 값이 오름차순으로 정렬된다.

distance 결과

node	[0]	[1]
1	0	∞
2	2	∞
3	6	7
4	4	5
5	6	∞

 

3️⃣ 4번 노드에서 갈 수 있는 노드 갱신

• 힙 큐 = [[4, 4], [5, 4], [6, 3], [7, 3], [6, 5]]
• pq에서 pop한 node = 4, cost = 4
• 노드 4에서 갈 수 있는 노드, 현재 cost(4) + 인접노드까지의 거리 정보 삽입 (None)
• 결과 : 힙 큐 = [[5, 4], [6, 3], [7, 3], [6, 5]]

distance 결과 (앞과 동일)

node	[0]	[1]
1	0	∞
2	2	∞
3	6	7
4	4	5
5	6	∞

... 반복 🔁

 

📍 전체 코드

# 다익스트라
import sys
import heapq  # 우선순위 큐

input = sys.stdin.readline
n, m, k = map(int, input().split())

direction = [[] for _ in range(n + 1)]
distance = [[sys.maxsize] * k for _ in range(n + 1)]

for _ in range(m):
    start, end, time = map(int, input().split())

    direction[start].append([end, time])

distance[1][0] = 0
pq = [(0, 1)]

while pq:
    cost, node = heapq.heappop(pq)
    
    for nextNode, nextCost in direction[node]:
        sumCost = cost + nextCost

        if distance[nextNode][k - 1] > sumCost:
            distance[nextNode][k - 1] = sumCost
            distance[nextNode].sort()

            heapq.heappush(pq, [sumCost, nextNode])


for dis in distance[1:]:
    if dis[k - 1] == sys.maxsize:
        print(-1)
    else:
        print(dis[k - 1])