[백준 BOJ_1504] 특정한 최단 경로 Python 풀이

출처: 백준 온라인 저지

문제

풀이

최단거리 알고리즘인 다익스트라(dijkstra)를 이용하여 구현해주었습니다.

heapq 모듈을 사용하여 구현해주었고, dist를 math.inf로 초기화해주고, 다익스트라를 통해 src에서 출발하여 도착할 수 있는 최단경로를 update해주었습니다.

다익스트라를 다 돌고나면 dist에는 경로가 있는 한 최단거리로 update가 되니, 각각 1, v1, v2의 dist를 구해주었습니다. 그 후에는 아래의 두 가지 경우를 고려해주었습니다.

• 1 → v1 → v2 → N
• 1 → v2 → v1 → N

v1을 먼저 방문하고 v2를 방문한 뒤 N으로 가거나, v2를 먼저 방문하고 v1을 방문한 뒤 N으로 가는 두 가지 경우 중 더 짧은 거리를 answer에 저장해줍니다.

마지막에는 answer를 출력해줍니다. answer가 만약 math.inf라면 경로가 없는 경우이기 때문에 -1를 출력해줍니다.

코드

import math
from collections import defaultdict
from heapq import heappush, heappop

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

adj = defaultdict(list)

for _ in range(E):
  a, b, c = map(int, input().split())
  heappush(adj[a], (c, b))
  heappush(adj[b], (c, a))

def dijkstra(src):
  dist = [ math.inf for _ in range(N+1)]
  dist[src] = 0
  hq = [(0, src)]

  while hq:
    w, v = heappop(hq)
    if w > dist[v]:
      continue
    for nw, nv in adj[v]:
      nw += w
      if dist[nv] > nw:
        dist[nv] = nw
        heappush(hq, (nw, nv))
  return dist

v1, v2 = map(int, input().split())
dist1 = dijkstra(1)
distv1 = dijkstra(v1)
distv2 = dijkstra(v2)
way1 = dist1[v1] + distv1[v2] + distv2[N]
way2 = dist1[v2] + distv2[v1] + distv1[N]

answer = min(way1, way2)
print(-1 if answer == math.inf else answer)