[백준] 1504번 특정한 최단 경로(다익스트라 알고리즘)(Python - 파이썬)

백준 1504번 문제입니다. (solved.ac)기준 골드 4 문제입니다.

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

1504번: 특정한 최단 경로

문제 접근

두 정점을 반드시 지나는 최단 경로의 길이를 출력해야하는 문제입니다.

다익스트라 알고리즘은 출발노드에서 다른 모든 노드까지의 최소거리를 구하는 알고리즘이기 때문에 다익스트라 알고리즘을 이용하여

1 -> v1 -> v2 -> n 혹은 1 -> v2 -> v1 -> n 두 가지 방법(양방향 길이 존재) 중 더 작은 방법을 출력하면 되는 문제입니다.

1 -> v1 -> v2 -> n 은 1에서 v1으로 가는 최단 경로 + v1에서 v2로 가는 최단 경로 + v2에서 n으로 가는 최단 경로를 구하면 되고

1 -> v2 -> v1 -> n 은 1에서 v2로 가는 최단 경로 + v2에서 v1으로 가는 최단 경로 + v1에서 n으로 가는 최단 경로를 구하면 됩니다.

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

    q = []
    heapq.heappush(q, (0, start))

    while q :

        dis, now = heapq.heappop(q)

        for i in graph[now] :
            cost = dis + i[1]

            if distance[i[0]] > cost:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

    return distance

위와 같은 형태의 다익스트라 알고리즘 함수를 구현하여 시작 지점(start)으로부터 다른 노드까지의 최단 경로(distance)를 반환해주도록 하였습니다. 

 

from_1 변수에 시작 지점이 1일 때 다른 노드까지의 최단 경로를 저장

from_v1 변수에 시작 지점이 v1일 때 다른 노드까지의 최단 경로를 저장

from_v2 변수에 시작 지점이 v2일 때 다른 노드까지의 최단 경로를 저장

코드로 나타내면 아래와 같습니다.

from_1 = dijkstra(1)
from_v1 = dijkstra(v1)
from_v2 = dijkstra(v2)

문제에 나와 있는 예제 1번 기준으로 아래와 각각 아래의 리스트가 저장됩니다.

from_1
[1000000000, 0, 3, 5, 4]

from_v1
[1000000000, 3, 0, 3, 4]

from_v2
[1000000000, 5, 3, 0, 1]

따라서 from_1[3]은 1에서 3으로 가는 최단 경로, from_v1[4]는 v1에서 4로 가는 최단 경로를 뜻합니다.

 

1 -> v1 -> v2 -> n 혹은 1 -> v2 -> v1 -> n 중 더 작은 값을 구해주면 되기 때문에

ans = min(from_1[v1] + from_v1[v2] + from_v2[n], from_1[v2] + from_v2[v1] + from_v1[n])

를 이용하여 ans에 더 작은 값을 저장합니다.

마지막으로 경로가 없다면 -1을 출력해야하기 때문에 조건에 맞게 출력해주면 됩니다!

if ans >= INF : print(-1)
else : print(ans)

정답 코드

import heapq
import sys
INF = int(1e9)

n, e = map(int, sys.stdin.readline().rstrip().split())

graph = [[] for _ in range(n+1)]

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

v1, v2 = map(int, sys.stdin.readline().rstrip().split())

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

    q = []
    heapq.heappush(q, (0, start))

    while q :

        dis, now = heapq.heappop(q)

        for i in graph[now] :
            cost = dis + i[1]

            if distance[i[0]] > cost:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

    return distance

from_1 = dijkstra(1)
from_v1 = dijkstra(v1)
from_v2 = dijkstra(v2)

ans = min(from_1[v1] + from_v1[v2] + from_v2[n], from_1[v2] + from_v2[v1] + from_v1[n])

if ans >= INF : print(-1)
else : print(ans)