[BOJ]1504 특정한 최단 경로 - 파이썬

https://www.acmicpc.net/problem/1504

풀이

이 문제는 무난한 다익스트라 알고리즘을 사용하는 문제였다.
단 s1,s2를 무조건 지나야 한다는 점을 염두에 두어야한다.

내가 생각한 방법은 시작점은 무조건 1 이므로
1->s2->s1->n
1->s1->s2->n
을 각각 계산하고 최솟값을 출력하는 방식이었다.
즉 다익스트라를 총 6번 실행했다. (s1->s2, s2->s1은 같기 때문에 5번으로도 가능하다)

import heapq

import math
from collections import defaultdict

n, m = map(int, input().split(' '))
visited = defaultdict(bool)
graph = defaultdict(list)
for i in range(m):
  a, b, c = map(int, input().split(' '))
  graph[a].append((b, c))
  graph[b].append((a, c))
s1, s2 = map(int, input().split(' '))


def dijkstra(start, end):
  distance = [math.inf]*(n+1)
  q = []
  heapq.heappush(q, (0, start))
  distance[1] = 0
  if start == end:
    return 0
  while q:
    dist, now = heapq.heappop(q)
    if distance[now] < dist:
      continue
    for next_node, value in graph[now]:
      cost = dist+value
      if cost < distance[next_node]:
        distance[next_node] = cost
        heapq.heappush(q, (cost, next_node))
  return distance[end]


total1 = dijkstra(1, s1) + dijkstra(s1, s2)+dijkstra(s2, n)
total2 = dijkstra(1, s2)+dijkstra(s2, s1)+dijkstra(s1, n)
print(-1 if min(total1, total2) == math.inf else min(total1, total2))