다익스트라 알고리즘을 사용하여 구현하면 된다.
처음에는 힙큐에 데이터를 넣어줄때 (정점, 비용) 순서로 넣어주는게 자연스러울것 같아서 파이썬의 클래스와 스페셜 메소드 lt를 사용하여 힙큐의 우선순위를 커스텀 하는 방식으로 구현하였다.
코드는 다음과 같다.
class node:
def __init__(self, v, w):
self.v = v
self.w = w
def __lt__(self, other):
return self.v < other.w # 오름차순
def dijkstra(start):
heap = []
distance = [INF] * (V + 1)
distance[start] = 0
heapq.heappush(heap, node(start, 0))
그러나 위와 같이 힙큐에 데이터를 넣어줄때마다 객체를 만들어서 넣어주게 되니 시간초과가 발생하였다.
결국 맘에는 안들지만 (비용, 정점) 순서대로 튜플 형태로 힙큐에 넣는 방식으로 구현하였고, 이렇게 구현하니 통과가 되었다.
비록 원하는 방향대로 구현하지는 못했지만 힙큐의 우선순위를 커스텀 해주는 방식은 이후 다른 알고리즘 문제를 해결할때 사용될 수 있다고 생각한다.
import sys
import heapq
input = sys.stdin.readline
INF = sys.maxsize
V,E = map(int, input().split(' '))
K = int(input())
graph = [[] for _ in range(V + 1)]
for i in range(E):
u,v,w = map(int, input().split(' '))
graph[u].append((v, w))
def dijkstra(start):
heap = []
distance = [INF] * (V + 1)
distance[start] = 0
heapq.heappush(heap, (0, start))
while heap:
weight, vertex = heapq.heappop(heap)
if weight > distance[vertex]: continue
for v, w in graph[vertex]:
if distance[v] > weight + w:
distance[v] = weight + w
heapq.heappush(heap, (weight + w, v))
for i in range(1, V+1):
print(distance[i]) if distance[i] != INF else print('INF')
dijkstra(K)