[백준] 1504 - 특정한 최단 경로 [Python(파이썬)]

문제

www.acmicpc.net/problem/1504

1504번: 특정한 최단 경로

풀이

백준 1753번 문제인 최단경로 문제와 유사한 문제다. 마찬가지로 다익스트라 알고리즘을 적용하여 해결할 수 있다.

다만, 이 문제 같은 경우 두 정점(v1, v2)을 반드시 거쳐야 한다는 제약조건을 고려해야 한다.

 

1번 노드부터 출발하여 n번 노드까지 이동하는 경로는 다음과 같이 두 가지로 나뉜다.

1) v1번 노드를 먼저 방문한 후, v2번 노드를 방문하는 경로: 1 - v1 - v2 - n(dst)

2) v2번 노드를 먼저 방문한 후, v1번 노드를 방문하는 경로: 1 - v2 - v1 - n(dst)

 

총 3번의 다익스트라 함수를 통해서 1번 경로와 2번 경로의 길이를 구할 수 있다.

첫 번째 다익스트라(1에서 출발): 1 ~ n의 최단경로를 구한다. → 1 ~ v1, 1 ~ v2의 길이를 구할 수 있다.

두 번째 다익스트라(v1에서 출발): v1 ~ n의 최단경로를 구한다. → v1 ~ n의 길이를 구할 수 있다.

세 번째 다익스트라(v2에서 출발): v2 ~ n의 최단경로를 구한다. → v2 ~ n의 길이를 구할 수 있다.

from_1 = dijkstra(1)
from_v1 = dijkstra(v1)
from_v2 = dijkstra(v2)
    
path1 = from_1[v1] + from_v1[v2] + from_v2[n]
path2 = from_1[v2] + from_v2[v1] + from_v1[n]

이때 경로가 존재하지 않으면 해당 경로는 INF를 값으로 가지게 된다. 만약 최솟값이 INF라면 -1을 반환한다.

 

코드

import sys, collections, heapq

INF = float('inf')

n, e = map(int, sys.stdin.readline().split())
g = collections.defaultdict(list)

for _ in range(e):
    a, b, c = map(int, sys.stdin.readline().split())
    g[a].append([b, c])
    g[b].append([a, c])

v1, v2 = map(int, sys.stdin.readline().split())
    
def dijkstra(s):
    dist = [INF] * (n+1)
    dist[s] = 0
    q = [[dist[s], s]]
    while q:
        u = heapq.heappop(q)[1]
        for v, c in g[u]:
            if dist[v] > dist[u] + c:
                dist[v] = dist[u] + c
                heapq.heappush(q, [dist[v], v])
    return dist

def solve():
    from_1 = dijkstra(1)
    from_v1 = dijkstra(v1)
    from_v2 = dijkstra(v2)

    path1 = from_1[v1] + from_v1[v2] + from_v2[n]
    path2 = from_1[v2] + from_v2[v1] + from_v1[n]

    result = min(path1, path2)
    
    if result < INF:
        return result
    else:
        return -1

print(solve())