백준 1854번: K번째 최단경로 찾기 (python)
by VICENTE97P4
Feb. 12, 2022, 3:01 p.m.
문제 링크: https://www.acmicpc.net/problem/1854
오랜만입니다.
자주 써야 하는데.. 방학 안에 스프링을 한 번 전체적으로 보려고 하다 보니 시간이 부족하네요..
똑같은 것만 계속 하다 보면 지루해지기 마련이죠..
네, 참지 못하고 알고리즘 문제를 풀어봤습니다.
문제에서 요구하는 바는 간단합니다.
1번 도시에서 각 도시까지의 K번째로 짧은 경로를 출력하는 것입니다.
짧은 경로라.. 최소 경로는 BFS, Dijkstra, floyd-warshall 등을 이용할 수 있겠고..
1번 도시에서 시작하니까 floyd-warshall은 과하고,
모든 도시까지의 최단거리니까 다익스트라가 좋아보입니다.
네, 이번 문제는 다익스트라를 응용한 문제입니다.
그런데 다익스트라는 최단거리는 구할 수 있는데 K번째로 짧은 거리는 어떻게 구할까요?
다익스트라는 각 노드까지의 최단거리를 저장하는 distance 배열이 핵심입니다.
뭔가 이 distance 배열에 각 도시를 통과하는 모든 경우를 다 저장해준 다음,
각 도시의 distance를 모두 sort해주고, K번째 거리를 출력하면 되지 않나!!
생각해봤지만 그러면 dijkstra 알고리즘이 끝나지 않겠죠?
그럼 어떡할까.. 도시를 방문할 때, K번째 시간보다 걸리는 시간이 더 길면 그냥 heap에 넣어주지 않으면 됩니다.
그럼 K번째 걸리는 시간을 어떻게 바로 알 수 있을까요?
일단 K번째 걸리는 시간을 알려면 K번째 앞에 걸리는 시간까지 모두 다 알아야 합니다.
왜냐하면 지금 당장 K번째 걸리는 시간을 구했다 하더라도, 나중에 이 도시를 방문하는 K번째보다 짧은 시간이 걸리는
경로가 들어올 수도 있으니까요
그럼 도시를 방문한 시간들을 K번째까지 저장해야 하는데,
그냥 리스트로 저장하면 한 번 시간 저장을 할 때마다 O(N)이 걸리게 됩니다.
예를 들어 [2, 8, 15, 30] 이런 리스트가 있을 때, 16이라는 수가 들어온다고 합시다.
그럼 16이 들어갈 위치를 찾고, 뒤에 있는 30을 한 칸 뒤로 밀어주고.. 그래서 N이 걸립니다.
그럼 뭐 어쩌란 말인가.. 어떤 자료구조를 쓰란 말인가.. 생각해보니
1. K번째 이전 시간들을 모두 저장해야 한다.
2. K번째에 걸리는 시간을 상수 시간에 알 수 있어야 한다.(K번째보다 더 긴 시간이 걸리면 그 경로는 그냥 무시해야 하므로)
3. 시간의 추가와 제거에 많은 시간이 걸리면 안된다.
이걸 모두 만족하는 자료구조는 priority queue입니다.
이전 수들은 모두 저장해야 하고, 제일 크거나 작은 값은 바로 알 수 있고, 추가와 제거에 시간이 많이 안 걸리고..
이 문제는 거리 저장에 priority queue를 이용하는 독특한 다익스트라 문제였습니다.
코드를 보겠습니다.
import sys;
input = sys.stdin.readline;
import heapq
n, m, k = map(int, input().split())
graph = [[] for _ in range(n + 1)]
for _ in range(m):
a, b, c = map(int, input().split())
graph[a].append([b, c])
distance = [[] for _ in range(n + 1)]
distance[1].append(0)
pq = [[0, 1]]
while pq:
cost, node = heapq.heappop(pq)
for nex, weight in graph[node]:
total = cost + weight
if len(distance[nex]) < k:
heapq.heappush(distance[nex], -total)
heapq.heappush(pq, [total, nex])
elif total < -distance[nex][0]:
heapq.heappop(distance[nex])
heapq.heappush(distance[nex], -total)
heapq.heappush(pq, [total, nex])
for i in range(1, n + 1):
if len(distance[i]) < k:
print(-1)
else:
print(-distance[i][0])
길어보이는데, 사실 다익스트라 로직만 안다면 추가된 부분은 별로 없습니다.
heap은 max-heap을 위해 heap에 들어가는 수는 음수로 바꿔줬습니다.
for nex, weight in graph[node]:
total = cost + weight
if len(distance[nex]) < k:
heapq.heappush(distance[nex], -total)
heapq.heappush(pq, [total, nex])
elif total < -distance[nex][0]:
heapq.heappop(distance[nex])
heapq.heappush(distance[nex], -total)
heapq.heappush(pq, [total, nex])
if 조건문에서 다음 방문할 도시에 저장된 시간들의 개수가 k개보다 작으면 무조건 그냥 저장합니다.
당연하죠, 시간이 얼마가 걸리던, K개를 일단 채우는 것이 우선이니까요.
elif 조건문에서 다음 방문할 도시에 저장된 시간들 중 제일 큰 수가 K번째 수입니다.
시간을 K개만 저장하기 때문에 당연한 부분이죠.
그래서 0번 index인 시간보다 작으면 제일 큰 수를 pop하고, 새로운 시간을 우선순위 큐에 넣습니다.
그래야 K개의 시간이 유지되겠죠.
이 방식대로라면 정확하게 K번째 걸리는 시간을 구할 수 있습니다.
for i in range(1, n + 1):
if len(distance[i]) < k:
print(-1)
else:
print(-distance[i][0])
정답을 출력하는 부분입니다.
K번째 시간이 없을 수 있으니까 distance에 저장된 시간이 k개보다 작으면 -1을,
k개 이상이면 가장 큰 수를 출력합니다.(max-heap이니까 0번 index가 가장 큰 수입니다.)
고민을 많이 해서 힘든 문제였지만, 훌륭한 배움을 얻은 문제였습니다.
이전 수를 모두 저장해야 할 때, 특정 번째의 수를 구할 때 우선순위 큐를 이용할 수도 있다는 교훈..
그런데 스프링 공부는 언제 다 할까요.. 문제를 풀어서 좋지만 남은 일이 많아 갑갑함이 남네요.
백준 dijkstra 0 view 1065