STUDY/Algorithm

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

sinawi95 2021. 4. 27. 21:42
728x90

www.acmicpc.net/problem/1504

 

1504번: 특정한 최단 경로

첫째 줄에 정점의 개수 N과 간선의 개수 E가 주어진다. (2 ≤ N ≤ 800, 0 ≤ E ≤ 200,000) 둘째 줄부터 E개의 줄에 걸쳐서 세 개의 정수 a, b, c가 주어지는데, a번 정점에서 b번 정점까지 양방향 길이 존

www.acmicpc.net

import sys
input = sys.stdin.readline
from heapq import heappop, heappush
N, E = map(int, input().split())
link_arr = [list() for _ in range(N + 1)]
for _ in range(E):
    n1, n2, w = map(int, input().split())
    link_arr[n1].append((n2, w))
    link_arr[n2].append((n1, w))
v1, v2 = map(int, input().split())

# inf = 800*1000 설정
# 모든 정점을 지나가도 지나가는 최소 정점의 개수는 (N-1)개 이기때문
result_list = []
def dijkstra(start, end1):
    if start==end1:
        return 0
    dist = [800000] * (N + 1)
    visit = [0] * (N + 1)
    q = [(0, start)] # weight, start
    dist[start] = 0
    while q:
        w, cur = heappop(q)
        visit[cur] = 1
        if cur == end1:
            return dist[cur]
        for i in range(len(link_arr[cur])):
            adj, adj_w = link_arr[cur][i]
            if not visit[adj] and dist[adj] > dist[cur] + adj_w:
                dist[adj] = dist[cur] + adj_w
                heappush(q, (dist[adj], adj))
    return 8000000
# 1. 1-> v1 -> v2 -> N
result1 = dijkstra(1, v1) + dijkstra(v1, v2) + dijkstra(v2, N)
# 2. 1-> v2 -> v1 -> N
result2 = dijkstra(1, v2) + dijkstra(v2, v1) + dijkstra(v1, N)
result = min(result1, result2)
if result >= 800000:
    result = -1

print(result)

다익스트라를 활용한 최소 거리 문제이다

다익스트라를 사용하면서 end에 해당하는 정점에 오면 그 점까지의 거리를 반환하고 길이 없으면 최대값을 반환한다.

최대값은 800000으로 잡아놓았는데 모든 정점을 지나가도 지나가는 최소 정점의 개수는 (N-1)개 이기때문에 (정점의 개수의 최대값) * (가중치의 최대값)으로 설정하였다

v1, v2를 꼭 지나야하기때문에 1에서 시작해서 v1, v2를 거쳐 N을 가는 것과 1에서 시작해서 v2, v1을 거쳐 N을 가는 것, 두가지 선택지가 있다

두가지 방법중 더 짧은것을 구하고 그 값이 inf=800000 보다 작으면 값을 출력하고 아니면 -1을 출력하면 된다.

다익스트라만 알면 되는 문제이므로 조금 쉽게 풀었다.