[PS] BOJ 1753 / 최단경로

[PS] BOJ 1753 / 최단경로
문제 링크: https://www.acmicpc.net/problem/1753
Thumbnail: Photo by oxana v (Unsplash)

정말 정석적인 다익스트라 구현 문제입니다.

풀이

다익스트라(Dijkstra)!

다익스트라 알고리즘은 그리디하게 현재 방문한 정점들로부터 가장 가까운 거리에 위치하는 정점을 찾아가는 방식입니다. 이번 문제는 정점 \(K\)에서 나머지 다른 정점들까지의 최단 거리를 계산해야 하니, 다익스트라 알고리즘이 적합합니다.

다익스트라 알고리즘의 구현에 앞서, 먼저 그래프를 구현해야 합니다. 정점은 최대 20,000개이나 간선은 최대 300,000이므로, 인접 리스트 방식으로 그래프를 구현했습니다.

문제에서 방향 그래프라고 언급했고, 각 간선의 가중치는 10 이하의 자연수로 주어짐을 유의해 그래프를 정의합니다.

INF = 100_000_001                   # 방문 불가능한 정점의 거리를 나타낼 상수
input = open(0).readline
V, E = map(int, input().split())    # 정점의 개수, 간선의 개수
K = int(input().strip())            # 시작 정점 번호
graph = [[] for _ in range(V + 1)]
for _ in range(E):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))

인접 리스트로 구현한 그래프

이제 다익스트라 알고리즘을 구현해야 합니다. 그냥 구현해도 되지만, 우선순위 큐를 사용하면 시간 복잡도를 개선할 수 있습니다. 파이썬에서는 PriorityQueue를 사용하면 됩니다.

💡 수정)
Python3의 queue 모듈은 멀티스레드 환경에서의 큐 처리를 위한 모듈입니다.
단순히 큐의 기능이 필요한 경우 collections모듈의 deque (double-ended queue)를 사용하는게 성능 면에서 유리합니다. 또한, 우선순위 큐 기능이 필요한 경우 heapq모듈을 활용해 list를 우선순위 큐로 사용할 수 있습니다. Python 환경에서 PS를 할 계획이시라면 이 두가지 방법이 적절합니다.

이 글은 해당 사실을 알기 전에 작성되어 PriorityQueue를 사용해 구현되었습니다. 이 글의 코드를 수정하진 않고, 대신 heapq를 활용한 다익스트라 구현을 다루는 다른 글을 참고해 주세요.
from queue import PriorityQueue

distance = [INF for _ in range(V + 1)]
distance[K] = 0
pq = PriorityQueue()  # 우선순위 큐 초기화
pq.put((0, K))
while not pq.empty():
    dist, current_node = pq.get()  # 현재 노드와 거리를 가져옴
    if distance[current_node] < dist:  # 이미 더 짧은 경로가 있는 경우
        continue
    for neighbor, weight in graph[current_node]:  # 현재 노드의 이웃 노드와 가중치 확인
        new_distance = dist + weight
        if new_distance < distance[neighbor]:  # 새로운 거리가 더 짧은 경우
            distance[neighbor] = new_distance
            pq.put((new_distance, neighbor))

다익스트라 알고리즘

이후 distance배열을 출력해주면 됩니다. print문의 호출을 줄이기 위해 단일 문자열로 만든 후 출력했습니다.

print("\n".join(map(lambda x: str(x) if x < INF else "INF", distance[1:])))

결과 출력

전체 코드

from queue import PriorityQueue

INF = 100_000_001
input = open(0).readline
V, E = map(int, input().split())    # 정점의 개수, 간선의 개수
K = int(input().strip())            # 시작 정점 번호
graph = [[] for _ in range(V + 1)]  # 그래프 초기화
for _ in range(E):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))              # u에서 v로 가는 간선과 가중치 추가 (주어진 그래프는 방향 그래프이므로 u->v로만 간선 추가)

distance = [INF for _ in range(V + 1)]
distance[K] = 0
pq = PriorityQueue()  # 우선순위 큐 초기화
pq.put((0, K))
while not pq.empty():
    dist, current_node = pq.get()  # 현재 노드와 거리를 가져옴
    if distance[current_node] < dist:  # 이미 더 짧은 경로가 있는 경우
        continue
    for neighbor, weight in graph[current_node]:  # 현재 노드의 이웃 노드와 가중치 확인
        new_distance = dist + weight
        if new_distance < distance[neighbor]:  # 새로운 거리가 더 짧은 경우
            distance[neighbor] = new_distance
            pq.put((new_distance, neighbor))

print("\n".join(map(lambda x: str(x) if x < INF else "INF", distance[1:])))

solution.py

시행착오

처음에, 가중치는 10 이하의 자연수로 부여된다는 글만 보고 최단거리의 INF를 11로 잡는 실수를 했었습니다... 이후 충분히 큰 수로 바꿔 해결했습니다. 😄