[백준/Python] 1504 특정한 최단 경로

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

알고리즘

출발점이 1 종료지점이 n이고  v1, v2를 지나는 최단경로의 길이를 출력하는 문제이다. 가능한 경로는1 -> v1 -> v2 -> n / 1 -> v2 -> v1 -> n 두가지 경로가 존재하므로 각각의 경우를 전부 계산해서 더해준 뒤에 더 작은 값을 출력해주면 된다.출발점을 1, v1, v2 로 하여 다익스트라를 진행한 뒤에 각각의 경우를 더해서 최단경로를 계산해주었다.모든 경우마다 다익스트라를 실행하면 시간이 많이 소요되기 때문에, 각각의 경우에 대해 1번씩만 수행한 뒤에 distance배열을 리턴값으로 받아서 마지막에 더해주었다.

코드

import sys
import heapq
input = sys.stdin.readline
INF = int(1e9)

n,e = map(int,input().split())
graph = [[] for _ in range(n+1)]

for _ 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())

def dijkstra(start):
    distance = [INF] * (n + 1)
    q = []
    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]))

    return distance

dist_from_1 = dijkstra(1)
dist_from_v1 = dijkstra(v1)
dist_from_v2 = dijkstra(v2)

way_1 = dist_from_1[v1] + dist_from_v1[v2] + dist_from_v2[n] # 1->v1->v2->n
way_2 = dist_from_1[v2] + dist_from_v2[v1] + dist_from_v1[n] # 1->v2->v1->n

ans = min(way_1, way_2)
print(ans if ans < INF else -1)