popeye0618

Backend Developer

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

ko 1 조회수 시리즈 · 알고리즘 #알고리즘

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

골드 4 문제이며 제한 시간은 1초이다.

해결 과정

문제를 읽어보니 가중치가 있는 그래프에서의 최단거리 문제이므로 다익스트라 알고리즘을 사용하는 문제라고 생각이 들었다.

문제 조건에서 한 번 이동했던 간선도 다시 이동할 수 있다는 조건을 보고 다익스트라를 여러번 써서 목표 지점들 간의 최단 거리를 찾으면 되겠구나 싶었다.

따라서 start -> a -> b -> end 경우와 start -> b -> a -> end 경우를 비교해보았고, 정답 판정을 받았다.

코드

py
import sys, heapq
input = sys.stdin.readline

def dijkstra(start):
    hq = []
    visited = [float('inf')] * (n + 1)
    heapq.heappush(hq, (0, start))
    visited[start] = 0

    while hq:
        dist, now = heapq.heappop(hq)

        if dist > visited[now]:
            continue

        for next_node, weight in graph[now]:
            cost = dist + weight
            if cost < visited[next_node]:
                visited[next_node] = cost
                heapq.heappush(hq, (cost, next_node))

    return visited

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))

a, b = map(int, input().split())

from_start = dijkstra(1)
from_a = dijkstra(a)
from_end = dijkstra(n)

start_to_a = from_start[a]
start_to_b = from_start[b]
a_to_b = from_a[b]
a_to_end = from_end[a]
b_to_end = from_end[b]

answer = -1

# start -> a -> b -> end
path_1 = start_to_a + a_to_b + b_to_end
# start -> b -> a -> end
path_2 = start_to_b + a_to_b + a_to_end

if not (path_1 == float('inf') and path_2 == float('inf')):
    answer = min(path_1, path_2)

print(answer)

코드 설명

다익스트라를 세 번 사용해서 시작 지점이 start 인 경우, v1인 경우, end인 경우의 각 노드까지의 최단거리를 구해놓았다.

그 이후 start부터 v1, v2까지의 거리와, v1부터 v2까지의 거리, v1, v2,부터 end까지의 거리를 가지고 start -> a -> b -> end 경우와 start -> b -> a -> end 경우를 비교해서 작은 값을 출력했다.

그리고 두 경우가 둘 다 최대거리 float('inf') 인 경우 start -> end 까지 갈 수 없는 경우이므로 -1을 출력해준다.

댓글

0

아직 댓글이 없습니다.

댓글 쓰기