[Baekjoon/python] 특정한 최단 경로 #1504

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

1504번: 특정한 최단 경로

문제

방향성이 없는 그래프가 주어진다. 세준이는 1번 정점에서 N번 정점으로 최단 거리로 이동하려고 한다. 또한 세준이는 두 가지 조건을 만족하면서 이동하는 특정한 최단 경로를 구하고 싶은데, 그것은 바로 임의로 주어진 두 정점은 반드시 통과해야 한다는 것이다.

세준이는 한번 이동했던 정점은 물론, 한번 이동했던 간선도 다시 이동할 수 있다. 하지만 반드시 최단 경로로 이동해야 한다는 사실에 주의하라. 1번 정점에서 N번 정점으로 이동할 때, 주어진 두 정점을 반드시 거치면서 최단 경로로 이동하는 프로그램을 작성하시오.

코드

• 최단 경로? 다익스트라 알고리즘 적용
• 두 정점(v1, v2)를 반드시 거쳐야하니 "1 -> v1 -> v2 -> N" 과 "1 -> v2 -> v1 -> N" 두 가지 경로의 최단 거리를 구함

from collections import defaultdict
import sys
import heapq

input = sys.stdin.readline

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

graph = defaultdict(list)
for _ in range(E):
    a, b, c = map(int, input().split())
    graph[a].append((c, b))
    graph[b].append((c, a))

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

def dijkstra(src, dst):
    heap = [(0, src)]
    visited = set()
    
    while heap:
        cur_dist, cur_node = heapq.heappop(heap)
        if cur_node in visited:
            continue
        visited.add(cur_node)
        if cur_node == dst:
            return cur_dist
        for nxt_dist, nxt_node in graph[cur_node]:
            dist = cur_dist + nxt_dist
            if nxt_node not in visited:
                heapq.heappush(heap, (dist, nxt_node))

    return float('inf')

# 1 -> v1 -> v2 -> N
path1 = dijkstra(1, v1) + dijkstra(v1, v2) + dijkstra(v2, N)
# 1 -> v2 -> v1 -> N
path2 = dijkstra(1, v2) + dijkstra(v2, v1) + dijkstra(v1, N)
if float('inf') in [path1, path2]:
    print(-1)
else:
    print(min(path1, path2))