[BOJ, Python] 1753번_최단경로

이 글은 Velog에서 이전한 글입니다. Velog 원문 보기

백준 1753번

문제 제목을 보자마자 직감했다. "아 이거 다익스트라 알고리즘 문제구나"

방향그래프가 존재하며, 최단 경로를 묻는 문제는 높은 확률로 다익스트라 알고리즘을 묻는 문제이다.

파이썬에서 다익스트라 알고리즘을 구현하는 여러 방법이 있는데 나는 일반적인 우선 순위 큐를 사용해서 문제를 풀었다.

전체 풀이

import sys
input = lambda: sys.stdin.readline().rstrip()
import heapq

# 최단 경로 -> 우선 순위 큐, 다익스트라 알고리즘

def dijkstra(start_point):
    dist = [float('INF') for _ in range(V)] # 각 정점까지의 최단 경로를 담을 리스트
    dist[start_point] = 0
    pq = [(0, start_point)]

    while pq:
        current_dist, n = heapq.heappop(pq)

        if current_dist > dist[n]:
            continue

        for (v, d) in graph[n]:
            distance = d + current_dist
            if dist[v] > distance:
                dist[v] = distance
                heapq.heappush(pq, (distance, v))

    return dist

V, E = map(int, input().split())

start_point = int(input())-1

graph = [[] for _ in range(V)]

for _ in range(E):
    u,v,w = map(int, input().split())
    graph[u-1].append((v-1,w))

dist = dijkstra(start_point)

for d in dist:
    if d == float("INF"):
        print("INF")
    else:
        print(d)

그래프 문제를 풀 때 1-based 인덱스와 0-based 인덱스 중에서 편한 방법을 사용하면 되는데 나는 0-based 인덱스로 했다.

0-based Index?
단순하게 그래프를 list, dict 등으로 표현할 때 첫번째 지점을 0으로 한다는 것이다.
그래프 문제에서 정점이 주어질 때 보통 1부터 주어진다. 이를 정직하게
graph = [[] for _ in range(V+1)] 으로 리스트를 생성해서 0번 인덱스 자리는 버리는게 1-based Index이다.
반면 graph = [[] for _ in range(V)]으로 생성해서 정점 1을 0번 index에 배치하는게 0-based Index이다.

문제를 간단히 보면

def dijkstra(start_point):
    dist = [float('INF') for _ in range(V)] # 각 정점까지의 최단 경로를 담을 리스트
    dist[start_point] = 0
    pq = [(0, start_point)]

    while pq:
        current_dist, n = heapq.heappop(pq)

        if current_dist > dist[n]:
            continue

        for (v, d) in graph[n]:
            distance = d + current_dist
            if dist[v] > distance:
                dist[v] = distance
                heapq.heappush(pq, (distance, v))

    return dist

각 정점까지의 최단 경로를 담을 dist 리스트를 만들고 문제 요구 사항에 따라 시작 지점의 값은 0으로 한다.

pq는 (정점까지의 거리, 정점)이다. 여기서 인자의 순서도 중요한데 heapq를 사용해서 push할 때는 첫번째 인자 값을 기준으로 하기 때문에 최단 경로 문제에서는 첫 번째 인자를 '정점까지의 거리'로 해야한다.

이후는 간단하게 graph에서 정점을 꺼내오고 dist에 저장되어 있는 값보다 거리가 작다면 dist의 값을 갱신해준다.

결과