(Python) - BOJ(1854번) : K번째 최단경로 찾기

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

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

---

• 문제 : 

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

 

 

 

---

• 풀이 :

다익스트라를 이용한 최단거리 문제였습니다.

 

아이디어는 기본 다익스트라와 동일하게 i -> j로 가는 경로들에 대해 heap을 사용하여 짧은가중치부터 확인하며 최단거리로 갱신해줍니다. 이 때, K번째 수를 구하려면 K 이하의 모든 누적가중치를 알아야 하므로 i->j로 가는 배열의 크기를 (N+1) x (K) 만큼의 2차원 배열로 사용해줍니다.

또한, 현재 가려는 거리가 최단거리가 아니더라도 distance 배열에 넣어주어야 합니다. 우리는 가장 짧은 거리가 아닌 K번째 짧은 거리가 필요하니까요.

그러나, INF로 초기화된 가중치 배열에서 i번째 정점으로 가는 가중치인 distance[i]의 K-1번째 인덱스보다  현재 탐색하려는 정점의 가중치가 크거나 같은 경우 ( 아래의 코드 ) 

if distance[i][K-1] > w:

해당 거리는 탐색할 필요가 없습니다.

또한 distance 배열의 K-1번째에 가중치를 넣은 후에는 정렬을 한 번 시켜주어 K번째 작은 수가 인덱스 K-1에 오도록 하였습니다.

 

모든 탐색을 마친 후 distance의 [i][K-1] 의 값이 INF가 아니면 해당 값을 출력하고 INF면 -1을 출력합니다 (경로가 없음)

 

어러운 문제였습니다. 아이디어도 꽤나 어려웠고 가장 큰 문제는 다익스트라를 구현할 때 heap에 [가중치, 정점번호] 의 순으로 넣어줬어야 하는데 그 반대로 넣어주어 구조 자체가 망가져버려서 계속 시간초과가 떴던 문제였습니다.

시간초과의 늪에 빠졌습니다..

 

---

• 소스코드 : 

import sys
from heapq import heappush, heappop

input = sys.stdin.readline
inf = sys.maxsize

N,M,K = map(int,input().split())
weight = [[] for _ in range(N+1)]
distance = [[inf]*K for _ in range(N+1)]
heap = []

def dijkstra(start):
    heappush(heap,[0,start])
    distance[start][0] = 0
    while heap:
        cost,curr = heappop(heap)
        for i,c in weight[curr]:
            w = cost + c
            if w<distance[i][K-1]:
                distance[i][K - 1] = w
                distance[i].sort()
                heappush(heap,[w,i])
for i in range(M):
    a,b,c = map(int,input().split())
    weight[a].append([b,c])

dijkstra(1)

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