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

문제

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

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

문제 해결 전략

편의를 위해 시작하는 정점을 $S$ , 도착하는 정점을 $E$ , 반드시 지나야 하는 두 정점을 각각 $V_1$ , $V_2$ 라고 하자.

$S$ 에서 $E$ 로 가는 경로가 $V_1$ 과 $V_2$ 를 반드시 지난다고 했을때 경로를 나눠서 생각한다면 두 가지 경우가 있을 것이다.

$S → V_1 → V_2 → E$ : $V_1$ 을 먼저 방문한 후 $V_2$ 을 방문한다.
$S → V_2 → V_1 → E$ : $V_2$ 을 먼저 방문한 후 $V_1$ 을 방문한다.

따라서 두 가지 경로 중 더 가까운 경로가 최단 경로이며, 만약 두 경로 모두에서 경로 중 일부가 존재하지 않을 경우 최단 경로가 존재하지 않는 것으로 볼 수 있다.

다익스트라 알고리즘을 이용해 두 경로를 이루는 부분 경로를 구해보자.

문제의 그래프는 양방향 그래프로 서로 다른 두 정점간의 가는 최단경로와 오는 최단경로는 동일하다. 따라서 한 방향의 최단경로를 구함으로서 다른 방향의 최단경로의 길이도 알 수 있다. 위 두 경로를 이루는 6개 부분 경로를 살펴보면 모두 정점 $V_1$ 또는 $V_2$ 에서 출발하거나 도착한다. 따라서 다익스트라 알고리즘을 이용해 정점 $V_1$ , $V_2$ 에서 출발하는 최단경로를 각각 구하면 두 경로의 길이를 도출할 수 있다.

정답 코드

import sys
from queue import PriorityQueue


INF = 10 ** 8


def dijkstra(n, graph, s):
    distance = [INF for _ in range(n)]
    distance[s] = 0
    q = PriorityQueue()
    q.put((0, s))

    while not q.empty():
        _, v_start = q.get()
        for v_end in range(n):
            if v_end == v_start:
                continue
            w_end = graph[v_start][v_end]
            w_total = distance[v_start] + w_end
            if w_total < distance[v_end]:
                distance[v_end] = w_total
                q.put((w_total, v_end))

    return distance


def solution(n, graph, u, v):
    distance_u = dijkstra(n, graph, u)
    distance_v = dijkstra(n, graph, v)
    path_1 = distance_u[0] + distance_u[v] + distance_v[n - 1]
    path_2 = distance_v[0] + distance_u[v] + distance_u[n - 1]
    min_path = min(path_1, path_2)
    if min_path >= INF:
        return -1
    return min_path


n, e = [int(i) for i in sys.stdin.readline().rstrip().rsplit()]
graph = [[0 if i == j else INF for i in range(n)] for j in range(n)]
for _ in range(e):
    a, b, c = [int(i) for i in sys.stdin.readline().rstrip().rsplit()]
    a -= 1
    b -= 1
    graph[a][b] = graph[b][a] = c
u, v = [int(i) - 1 for i in sys.stdin.readline().rstrip().rsplit()]

print(solution(n, graph, u, v))