백준 1504번: 특정한 최단 경로 [python]

난이도 : 골드 4
유형 : 최단경로
https://www.acmicpc.net/problem/1504

문제 접근

예제 1을 그림으로 표현했다.
위와같이 1부터 N까지 노드가 있고 각 노드 사이에는 양방향 가중치 간선이 있다.
v1과 v2를 입력받은 후
1번부터 N번까지의 최단경로를 구할건데 반드시 v1,v2 이 두 노드를 거쳐가는 최단경로를 구해야 한다.
예제 1의 경우는 1->2->3->4 를 통해 7의 값이 답이 된다.

핵심 아이디어:
경로는 두가지로 좁힐 수 있다.
1. 1-> v1 -> v2 -> N
2. 1-> v2 -> v1 -> N

특정 노드에서 모든 경로까지의 최단거리는 다익스트라 알고리즘을 통해 구할수 있으므로 다익스트라 알고리즘을 세번 돌려서
1번에서 모든 정점까지 거리
v1에서 모든 정점까지 거리
v2에서 모든 정점까지 거리
를 구한 뒤
1, 2 중 최소거리를 구하면 된다.

해답 및 풀이

import heapq
import sys
input = sys.stdin.readline

INF = int(1e9)

# 1. 입력
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())

# 2. 다익스트라 함수

def dijkstra(start):
    distance = [INF] * (n + 1)
    distance[start] = 0
    heap = [(0, start)]

    while heap:
        dist, now = heapq.heappop(heap)
        if dist > distance[now]:
            continue
        for next_node, cost in graph[now]:
            new_cost = dist + cost
            if new_cost < distance[next_node]:
                distance[next_node] = new_cost
                heapq.heappush(heap, (new_cost, next_node))

    return distance

# 3. 필요한 거리 구하기
dist_from_1 = dijkstra(1)
dist_from_v1 = dijkstra(v1)
dist_from_v2 = dijkstra(v2)

# 경로 계산
path1 = dist_from_1[v1] + dist_from_v1[v2] + dist_from_v2[n]
path2 = dist_from_1[v2] + dist_from_v2[v1] + dist_from_v1[n]

result = min(path1,path2)

if result < INF:
    print(result)
else:
    print(-1)