백준 1753번 : 최단경로 in Python

1753번: 최단경로

 

www.acmicpc.net

이 문제는 각 간선마다 가중치가 있을 경우 시작점으로부터 다른 모든 정점으로의 최단 경로를 구하는 문제이다. 

기존의 최단경로 문제의 경우 모든 간선의 가중치가 같았지만, 이 문제는 각 간선마다 가중치가 달라 BFS로 해결이 불가능하다. 그래서 이 문제는 BFS가 아닌 다른 알고리즘을 사용해야 하는데, 바로 Dijkstra의 알고리즘을 이용한다.

Dijkstra 알고리즘은 이 문제를 풀기 위해 고안된 알고리즘으로, 전체적인 알고리즘은 BFS와 비슷하나, 세부적인 사항에서 간선의 가중치를 따짐에 따라 BFS와 차이가 발생한다. 우선 최단거리를 저장하기 위한 배열을 만드는 것은 BFS와 비슷하나, BFS를 보조하는 queue를 만드는 대신 priority queue를 만들어 queue에서 정점을 꺼낼 때 queue에 가장 먼저 들어온 정점을 꺼내는 게 아니라 queue에 있는 정점들 중에서 그 정점까지의 최단거리가 가장 짧은 정점을 꺼낸다. 이는 BFS와 달리 어떤 정점과 인접한 정점들을 탐색할 때 탐색하지 않은 정점들뿐만 아니라 탐색했던 정점들도 다시 탐색하기 때문이다. 

이는 단순히 간선의 개수가 곧 최단거리인 BFS와는 달리 경로에 있는 간선의 개수가 더 많아도 가중치를 고려한 최단거리가 더 짧은 경우가 나오는 경우가 충분히 가능하기 때문이다. 그런데 만약 단순히 후입선출을 통해 정점을 뽑는다면 아직 완전히 최적화되지 않은 정점에 대해 그 정점과 인접한 정점들에 대해 최단거리를 구하는 과정을 거치게 돼 온전히 최적화된 결과를 얻을 수 없게 된다.

예를 들어 정점 A에서 출발해 그 정점과 인접한 정점들 B, C에 대해 최단거리를 구했을 때 각각 15, 5가 나왔다고 하자. 그리고 정점 B를 먼저 순회했다고 하자. 만약 그냥 queue에 정점들을 저장했다면 정점 A에 대한 순회를 마치고 정점 B에 대해 그와 인접한 점들에 대해 순회하며 최단거리를 구했을 것이다. 그런데 정점 C부터 먼저 queue에서 꺼내고, 점 C에서 점 B로 가는 간선의 가중치가 3이었다면, 실제 정점 B까지의 최단거리는 15가 아니고 8이 됐을 것이다. 그러면 앞에서 최단거리가 15라고 생각하고 그와 인접한 정점들에 대해 최단거리를 구한 것보다, 최단거리가 8이라고 생각하고 그와 인접한 정점들에 대해 최단거리를 구한 것이 훨씬 짧을 것이다.

그래서 이 문제에서는 priority queue를 만들어 queue에 있는 정점들 중에서 그 정점까지의 최단거리가 가장 짧은 정점을 꺼내야 한다. 그리고 물론 이 정점에 대해 이미 인접한 정점들에 대한 순회가 끝난 경우도 있을 것이므로 이를 점검해 그러한 경우 건너뛴다. 이후 이 정점을 방문했다고 표시한 뒤, 이 정점과 인접한 정점들에 대해 순회하며 각 정점들의 최단거리를 계산한다. 이를 계산할 때에는 기존의 인접한 정점까지의 최단거리와 이 정점을 지나서 인접한 정점까지 가는 거리와 비교해 후자가 더 짧을 경우 기존의 최단거리를 줄이고 priority queue에 인접한 정점을 추가한다. 이러한 과정을 relax라고 한다.

이를 priority queue의 길이가 0이 될 때까지 계속 반복하면 시작점과 연결된 모든 점들에 대해 최단거리를 구할 수 있게 된다. 한편 최단거리를 저장하는 배열은 처음에 모든 값들을 sys.maxsize, 즉 가장 큰 값으로 한 뒤에 시작점에 해당하는 값을 0으로 지정한 뒤 시작한다. 그렇게 함으로써 마지막에 만약 모든 그래프를 순회했는데도 정점까지의 최단거리가 sys.maxsize면 해당 정점과 시작점이 연결되지 않았음을 의미하므로 INF를 출력하면 된다.

위 내용을 코드로 구현하면 다음과 같다.

import sys
import heapq
V, E = map(int, input().split())
K = int(input())
graph = [[] for _ in range(V+1)]
# 그래프를 저장할 때 정점과 간선의 가중치를 같이 저장한다.
for _ in range(E):
    u, v, w = map(int, input().split())
    graph[u].append([v, w])
# 최단거리를 저장하는 배열
dijk = [sys.maxsize for _ in range(V+1)]
# 해당 정점을 방문했었는지 확인하는 배열
visit = [False for _ in range(V+1)]
dijk[K] = 0
# 최단거리가 짧은 순으로 정점들을 저장하는 priority queue
heap = []
heapq.heappush(heap, [0, K])
while len(heap) > 0:
    # 최단거리가 가장 짧은 정점을 꺼낸다. 
    [d, u] = heapq.heappop(heap)
    if visit[u] == True: continue
    visit[u] = True
    for [v, w] in graph[u]:
        # relax 과정
        if dijk[v] > dijk[u] + w:
            dijk[v] = dijk[u] + w
            heapq.heappush(heap, [dijk[v], v])
for i in range(1, V+1):
    # 최단거리가 sys.maxsize인 경우 해당 정점이 시작점과 연결되지 않았음을 의미한다.
    if dijk[i] == sys.maxsize:
        print("INF")
    else:
        print(dijk[i])