[백준] 1504번 특정한 최단 경로 - 파이썬/다익스트라

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

import heapq
import sys
input = sys.stdin.readline

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

graph = [ [] for _ in range(N+1) ]
INF = 10**9

distance = [INF] * (N+1)

for i in range(E):
    a,b,c = map(int,input().split())

    graph[a].append((b,c))

    graph[b].append((a,c))


v1,v2 = map(int,input().split())

q = []

def dij(start):

    heapq.heappush(q,(0,start))
    distance[start] = 0
    while q:

        dist,now = heapq.heappop(q)

        if distance[now] < dist:
            continue

        for i in graph[now]:

            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q,(cost,i[0]))

ans1 = 0
ans2 = 0

dij(1)

ans1+=distance[v1]
ans2+=distance[v2]

distance = [INF] * (N+1)

dij(v1)

ans1+=distance[v2]
ans2+=distance[N]

distance = [INF] * (N+1)

dij(v2)

ans1+=distance[N]
ans2+=distance[v1]

print(ans1,ans2)

ans = min(ans1,ans2)
if ans >= INF:
    print(-1)
else:
    print(ans)

다익스트라를 여러번 해서 풀었다.

경로해야할 정점으로 v1 v2가 있으므로

1 -> v1 -> v2 -> N
1 -> v2 -> v1 -> N

이러한 두 경로가 있을 수 있다.

따라서 이 두 경로의 최단거리를 구한뒤 최소값을 출력했다.

만약, 정점끼리 연결되어있지 않다면 경로의 길이가 처음에 설정한 INF 이므로 이 경우 -1을 출력한다.