백준 1753번: 최단경로

이 문제도 다익스트라 알고리즘의 정석인 문제이다.

각 노드별 최단 거리를 구한 리스트를 리턴하여 print해주면 되기 때문이다.

import heapq

v, e = map(int, input().split())
k = int(input())
graph = [[] for _ in range(v + 1)]
INF = 1e8
distance = [INF] * (v + 1)
for i in range(e):
    start, end, w = map(int, input().split())
    graph[start].append((w,end))


def dijkstra(k):
    distance[k] = 0
    q = []
    heapq.heappush(q, (0,k))
    while q:
        dist,now = heapq.heappop(q)
        if dist > distance[now]:
            continue

        for a, b in graph[now]:
            if dist+a<distance[b]:
                distance[b]=dist+a
                heapq.heappush(q,(dist+a,b))

    return distance


lst = dijkstra(k)
for i in lst[1:]:
    if i == 1e8:
        print("INF")
    else:
        print(i)

노드의 갯수+1 만큼 distance에 [INF]요소를 추가해주고, graph에도 []요소를 추가해준다.

이제 출발 노드, 도착노드, 가중치를 입력받는데, 입력할때마다  graph[시작노드]에 도착노드, 가중치를 추가해준다.

이제 시작 노드를 구하고, 시작노드는 거리가 0이므로 시작노드와 0을 우선순위 큐에 추가해준다.

최단 거리인 0은 distance[시작노드]에 값인 INF에서 0으로 갱신해줘야한다.

이제 q에 요소가 없을때까지 반복문을 돌려준다.

q에서 pop를 하여 나온 거리와 노드를  dist,now에 저장해준다.

이때 이미 distance[now]의 값이 pop한 거리보다 작다면 최단 거리로만 이동해줘야하니깐 continue하여 그 거리를 이용하지 않게끔 해준다.

이제 graph[now]로 그 노드로부터 도착하는 노드와 거리를 탐색해준다.

distance[now]의 거리와 도착 노드까지의 거리를 합한값이 distance[도착노드]보다 작다면 작은 값으로 갱신을 해준다.

이제 그 노드와 거리를 q에넣어줘야한다.

반복문을 계속 돌려 distance 리스트를 리턴해주고, 리스트를 출력해줄때, 값이 1e8로 되어있는 요소는 INF로 출력되게끔 해준다. 그리고 1번 노드부터 출력해야 하므로 distance[1:]로 반복문을 돌려준다.

 

 

사실 한번 틀렸었다.

구글링을 해도 비슷하게 사람들이 코드를 짠거 같아서 반례들을 계속 찾아봤다.

틀린 이유는 시간 초과 였었는데, 그 이유는

튜플 간의 대소 관계는 앞 원소가 큰 것이 더 크고, 앞 원소가 같으면 뒤 원소가 큰 것이 더 크다. 그러니 (거리, 정점) 순서로 넣어야 다익스트라 알고리즘에 맞게 작동한다.

나는 정점 거리별로 넣어 정점은 똑같으니 그 뒤 원소를 비교하는 연산까지 했으니 시간 초과가 뜬것이다.