[백준 / 파이썬] 1504번: 특정한 최단 경로

1504번: 특정한 최단 경로

정답 코드

import sys
from heapq import heappush, heappop
input = sys.stdin.readline


def dijkstra(start):
    distance[start] = 0
    heap = [(0, start)]

    while heap:

        min_dist, min_node = heappop(heap)

        if min_dist <= distance[min_node]:

            for next_node, dist in graph[min_node]:
                new_dist = min_dist + dist
                if new_dist < distance[next_node]:
                    distance[next_node] = new_dist
                    heappush(heap, (new_dist, next_node))


v, e = map(int, input().split())
graph = [[] for _ in range(v + 1)]

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

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

INF = 999999999
for i in range(3):
    distance = [INF] * (v + 1)
    if i == 0:
        dijkstra(1)
        one_v1 = distance[v1]
        one_v2 = distance[v2]
    elif i == 1:
        dijkstra(v1)
        v1_v2 = distance[v2]
        v1_n = distance[v]
    else:
        dijkstra(v2)
        v2_v1 = distance[v1]
        v2_n = distance[v]

answer = min((one_v1 + v1_v2 + v2_n), (one_v2 + v2_v1 + v1_n))

if answer >= INF:
    answer = -1

print(answer)

문제 풀이

그 다음으로 dijkstra 함수를 정의합니다. 이 함수는 다익스트라 알고리즘을 사용하여 최단 경로를 계산하는 역할을 합니다. 다익스트라 알고리즘은 시작 노드로부터의 최단 경로를 찾는 알고리즘으로, 최소 힙을 사용하여 노드를 방문하는 순서를 관리합니다.


다익스트라 알고리즘을 간단히 설명하면 다음과 같습니다.

1. 시작 노드의 최단 거리를 0으로 초기화하고 최소 힙에 삽입합니다.
2. 최소 힙이 비어있지 않은 동안 다음을 반복합니다:
   최소 힙에서 최소 거리를 가진 노드를 꺼냅니다.
   해당 노드와 연결된 노드들에 대해 현재까지의 거리와 비교하여 더 짧은 거리로 갱신합니다.
   갱신된 노드들을 최소 힙에 삽입합니다.
3. 모든 노드에 대한 최단 거리가 계산되면 결과를 반환합니다.

주어진 코드에서는 dijkstra 함수 내에서 위 알고리즘을 구현하고 있습니다.


다음으로, 변수 v와 e에 입력 값을 할당합니다. 이는 각각 노드의 개수와 간선의 개수를 의미합니다. 그리고 graph라는 리스트를 생성하여 인덱스에 해당하는 노드와 연결된 노드들의 정보를 저장합니다. 간선의 정보는 입력을 받으면서 해당 노드의 리스트에 추가됩니다.


그 후에 v1과 v2를 입력 받습니다. 이는 반드시 거쳐야 하는 두 개의 노드를 의미합니다.


그 다음으로 INF라는 변수를 설정하여 초기 최단 거리를 무한대로 설정합니다. 이후 반복문을 통해 각 경우의 최단 거리를 계산합니다. 반복문은 0부터 2까지 총 세 번 실행되며, 각 경우에 따라 다익스트라 알고리즘을 적용합니다.

• 첫 번째 경우(i == 0): 시작 노드에서 v1과 v2를 각각 다익스트라 알고리즘을 적용하여 해당 노드까지의 최단 거리를 계산합니다.
• 두 번째 경우(i == 1): v1 노드에서 v2 노드까지의 최단 거리와 v1 노드에서 마지막 노드까지의 최단 거리를 계산합니다.
• 세 번째 경우(i == 2): v2 노드에서 v1 노드까지의 최단 거리와 v2 노드에서 마지막 노드까지의 최단 거리를 계산합니다.

마지막으로, 세 가지 경우 중 최단 거리의 합을 구하고 그 중 가장 작은 값을 answer 변수에 할당합니다. 만약 answer가 초기에 설정한 INF 값보다 크거나 같다면, 경로가 존재하지 않는 것이므로 -1을 할당합니다.


마지막으로, answer를 출력합니다.