백준 1504번 : 특정한 최단 경로 in Python

1504번: 특정한 최단 경로

 

www.acmicpc.net

이 문제는 이전의 가중치를 포함한 최단 경로 문제에서 특정 두 정점을 항상 지나야 한다는 조건이 붙은 문제이다.

만약 최단 경로에 조건에 있는 두 정점이 있는 경우에는 아무 문제가 없다. 그냥 정점 1에서 정점 N까지의 최단 경로의 길이를 출력하면 된다. 문제는 그렇지 않은 경우로, 이는 크게 두 가지 경우로 나눌 수 있다. 두 정점을 v1, v2라 할 때 v1을 먼저 갔다가 v2를 가고, 마지막으로 N으로 가는 방법과, v2를 먼저 갔다가 v1을 가고, 마지막으로 N으로 가는 방법이 있다. 물론 각 정점을 경유할 때마다 그 정점들 사이에서는 최단 경로로 움직여야 할 것이고, 따라서 각 정점마다 그 정점을 시작점으로 하는 Dijkstra 알고리즘을 실행해야 할 것이다.

이를 구현하기 위해 우선 Dijkstra 알고리즘을 시작점을 매개변수로 받는 함수로 만들고 시작점이 1일 때, v1일 때, 그리고 v2일 때 각각 Dijkstra 알고리즘 구현 함수를 실행해 그 결과 나온 각 경우에 대해 최단 거리를 저장한 배열을 얻는다. 그리고 이제 v1 -> v2 -> N 순서대로 경유했을 때 최단 거리와 v2 -> v1 -> N 순서대로 경유했을 때 최단 거리를 비교해 더 작은 값을 출력하면 된다. 단, 만약 v1, v2를 지나는 경로가 존재하지 않을 시 두 값 모두 sys.maxsize(INF)일 것이므로 이 경우에는 -1을 출력한다.

이를 코드로 구현하면 다음과 같다.

import sys
import heapq
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])
# 시작점을 매개변수로 받는 Dijkstra 알고리즘 구현 함수
def Dijkstra(v):
    global N
    dijk = [sys.maxsize for _ in range(N+1)]
    visit = [False for _ in range(N+1)]
    dijk[v] = 0
    heap = []
    heapq.heappush(heap, [0, v])
    while len(heap) > 0:
        [d, u] = heapq.heappop(heap)
        if visit[u] == True: continue
        visit[u] = True
        for [v, w] in graph[u]:
            if dijk[v] > dijk[u] + w:
                dijk[v] = dijk[u] + w
                heapq.heappush(heap, [dijk[v], v])
    # 시작점에서 출발했을 때 그래프의 각 정점까지의 최단 거리를 저장한 배열을 반환한다.
    return dijk
v1, v2 = map(int, input().split())
# 1, v1, v2에 대해 Dijkstra 알고리즘 함수를 실행한다.
dijk1 = Dijkstra(1)
dijk2 = Dijkstra(v1)
dijk3 = Dijkstra(v2)
# 1 -> v1 -> v2 -> N, 1 -> v2 -> v1 -> N 중 더 짧은 경로의 길이를 정한다.
res = min(dijk1[v1] + dijk2[v2] + dijk3[N], dijk1[v2] + dijk3[v1] + dijk2[N])
if res < sys.maxsize:
    print(res)
# v1과 v2를 경유해 N에 도달하는 경로가 없는 경우 -1을 출력한다.
else:
    print(-1)