[Python] 백준 1753번: 최단경로

www.acmicpc.net/problem/1753

1753번: 최단경로

import sys
import heapq
input = sys.stdin.readline


def bfs(start):
    global minWeights
    minWeights[start] = 0

    heap = []
    heapq.heappush(heap, (0, start))

    while heap:
        curWeight, curNode = heapq.heappop(heap)

        for adjacentNode, adjacentWeight in graph[curNode].items():
            nextAdjacentWeight = curWeight + adjacentWeight

            # 처음 방문이거나, 이미 방문하였으나 더 짧은 경로일 때
            if minWeights[adjacentNode] == -1 or minWeights[adjacentNode] > nextAdjacentWeight:
                minWeights[adjacentNode] = nextAdjacentWeight
                heapq.heappush(heap, (nextAdjacentWeight, adjacentNode))


# Input
V, E = map(int, input().split())
graph = {i: {} for i in range(1, V+1)}

K = int(input())
for _ in range(E):
    u, v, w = map(int, input().split())
    # u에서 v로 갈 때의 비용을 최소값으로 유지시킴
    graph[u][v] = min(graph[u][v], w) if graph[u].get(v) else w

# Solution
minWeights = [-1 for i in range(V+1)]
bfs(K)

for i in range(1, V+1):
    print(minWeights[i] if minWeights[i] != -1 else "INF")

 

우선, 입력을 받을 때, graph[u][v]에는 u에서 v로 이동할 때 드는 비용의 최소값이 저장되게 한다. 즉, 예시 케이스의 경우의 graph는 다음과 같다.

{1: {2: 2, 3: 3}, 2: {3: 4, 4: 5}, 3: {4: 6}, 4: {}, 5: {1: 1}}

 

아무튼, 이를 기반으로 bfs를 돌려주는데, 이 때 heapq 라이브러리를 사용하여 우선순위큐를 이용하여 힙을 사용하여야 한다. 처음에 이렇게 하지 않았는데, 시간 초과로 인해 문제가 풀리지 않았다. 힙의 각 원소에는 (무게, 노드) 순의 튜플이 저장된다.

힙의 원소를 하나 꺼내서, 그것의 인접 노드들을 for문으로 돌아준다. 이 때, 인접 노드들에 처음 방문(minWeight가 -1일 경우인데, minWeight가 처음에 -1로 초기화되기 때문)이거나, 그렇진 않지만 더 짧은 경로일 경우, 해당 인접 노드의 minWeight를 갱신시켜주고, heappush를 해준다.