백준 1854 <K번째 최단경로 찾기> Python
https://www.acmicpc.net/problem/1854
다익스트라 알고리즘 문제입니다. 사실 다익스트라 알고리즘보다는 우선순위 큐 문제라고 할 수 있을 것 같습니다.
다익스트라의 핵심 아이디어인 '다음 노드까지 도달 시간이 현재 최단시간보다 길면 큐에 삽입하지 않는다' 를 제외하기 때문입니다. 이는 당연히 k 번째 최단 시간을 구하기 위해서는 적어도 K 번은 같은 노드를 반복해서 방문해야 하기 때문입니다.
그렇게만 진행하면 우리는 크게 두 가지 해결해야 하는 과제가 생깁니다.
1. 만약 순환이 생긴다면 모든 노드를 큐에 집어넣었을 때 종료 조건이 존재하지 않습니다.
좋은 예시가 바로 예제입니다.
우리는 탐색을 진행하며 2 - 3 - 5 의 순환이 무한히 반복될 수 있음을 알 수 있습니다.
그렇다면 모든 노드에 K 번 이상 방문하면 종료하면 되지 않을까 생각을 해볼 수 있습니다. 하지만 1 에는 도달하는 엣지가 없어서 1 은 영원히 K (K>1) 번 이상 방문할 수 없습니다. 그렇다면 해당 방법도 종료 조건이 될 수는 없습니다.
2. K 번 이상 방문할 경우 최단 시간을 따로 저장해야 합니다.
단순히 모든 경우를 저장하면 K 가 커질수록 정렬에 부담이 갑니다. 결국 적당한 K 개만 골라 저장해야 합니다. 먼저 방문한 K 개만 저장할 경우 다익스트라는 그리디 알고리즘이기 때문에 후에 도달한 시간이 더 빠를 수 있으므로 적절하지 않습니다.
1 번과 2 번은 단순히 각각 풀 수 없고 함께 생각해주어야 합니다.
우선 2 번의 경우 우리는 최댓값 혹은 최솟값을 가장 빠르게 도출할 수 있는 자료구조를 알고 있습니다. 바로 우선순위 큐입니다. K 이하의 크기를 유지하는 우선순위 큐에 최단시간 후보들을 저장하면 입력 / 출력 / 저장 모두 효율적으로 진행할 수 있게 됩니다.
최단시간 후보를 우선순위 큐에 저장하기로 하고 1 번 문제를 생각해보겠습니다. 우리는 어떤 노드에 도달할 때 무조건 직전 노드를 통해 진입하게 됩니다. 그 말은 직전 노드를 K 번 방문했다면 어떤 노드에 진입할 기회가 K 번 있었다는 말이 됩니다.
좀 더 알고리즘에 가깝게 설명해보자면 직전 노드를 K 번 방문했다면 큐에 연결된 임의의 노드가 K 번 들어갑니다. 다익스트라 알고리즘에 의해 이미 방문했건 아직 큐에 들어있건 결국 미래에는 방문할 수 있게 됩니다. 즉 연결된 직전 노드를 K 번 방문했다면 다음 노드로 갈 수 있는 엣지는 이미 큐에 들어있으므로 직전 노드를 더 이상 방문하지 않아도 된다는 말이 됩니다. 물론 직전 노드로의 더 빠른 길을 찾게되면 다시 방문해야 합니다.
종합해보면 K 개의 후보 노드를 저장하는 우선순위 큐를 각 노드별로 운용합니다.
# k 번째 최단시간
kth_min_dists = [[] for _ in range(n+1)]
kth_min_dists[1].append(0)임의의 노드에 K 번 이상 방문했다면 해당 노드는 큐에 넣어지지 않아도 될 자격이 생기는데 해당 자격은 해당 노드로의 K 개의 후보 시간들보다 도달하는 시간이 길 때 충족됩니다. K 번 이상 방문했지만 조건에 충족되지 않는다면 K 개의 후보를 유지해야 하므로 후보 시간들 중 가장 느린 도달 시간을 제외해주고 현재 도달하는 시간을 넣어줍니다.
# 해당 노드까지 최단 시간 후보가 k 개 이하면
if len(nn_kth_nodes) < k:
# 후보에 넣어주기
hq.heappush(kth_min_dists[next_node], -next_time)
# 큐에 추가
hq.heappush(q, [next_time, next_node])
# k 개 존재하고 현재 k 번째 최단 시간보다 빠르면
elif len(nn_kth_nodes) == k and next_time < -nn_kth_nodes[0]:
# 후보 팝
hq.heappop(kth_min_dists[next_node])
# 후보에 넣어주기
hq.heappush(kth_min_dists[next_node], -next_time)
# 큐에 추가
hq.heappush(q, [next_time, next_node])import sys
import heapq as hq
input = sys.stdin.readline
def solution(n, m, k, edges):
# 그래프 정리
graph = [[] for _ in range(n+1)]
for a, b, c in edges:
graph[a].append([b, c])
# 큐
q = [[0, 1]]
# k 번째 최단시간
kth_min_dists = [[] for _ in range(n+1)]
kth_min_dists[1].append(0)
# 순회
while q:
# 소요 시간 / 현재 도시
now_time, now_node = hq.heappop(q)
# 순회
for next_node, time in graph[now_node]:
# 다음 노드까지 시간
next_time = now_time + time
# 다음 노드의 최단 시간 후보
nn_kth_nodes = kth_min_dists[next_node]
# 해당 노드까지 최단 시간 후보가 k 개 이하면
if len(nn_kth_nodes) < k:
# 후보에 넣어주기
hq.heappush(kth_min_dists[next_node], -next_time)
# 큐에 추가
hq.heappush(q, [next_time, next_node])
# k 개 존재하고 현재 k 번째 최단 시간보다 빠르면
elif len(nn_kth_nodes) == k and next_time < -nn_kth_nodes[0]:
# 후보 팝
hq.heappop(kth_min_dists[next_node])
# 후보에 넣어주기
hq.heappush(kth_min_dists[next_node], -next_time)
# 큐에 추가
hq.heappush(q, [next_time, next_node])
for i in range(1, n+1):
# k 번째 후보가 존재하면
if len(kth_min_dists[i]) == k:
print(-kth_min_dists[i][0])
else:
print(-1)
# 입력
n, m, k = map(int, input().strip().split())
edges = [list(map(int, input().strip().split())) for _ in range(m)]
solution(n, m, k, edges)