알고리즘 문제 풀이/Python

python 백준 23801 두 단계 최단 경로 2

맛대 2024. 11. 22. 16:21

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

로직

  • 다익스트라 알고리즘

    • 시작 지점에서 각 지점까지의 최소 거리를 구할 수 있음
  • X지점 - P개의 중간 정점 중 한 지점 - Z지점 으로 이동

    • X에서 다익스트라 = X부터 P에 속한 모든 지점 까지의 최단 거리 계산 가능
    • Z에서 다익스트라 = Z부터 P에 속한 모든 지점 까지의 최단 거리 계산 가능
    • 모든 P에 속한 지점(p)에 대해서 X->p->Z의 값 = X->p + Z->p
from heapq import heappush,heappop
import sys
input = sys.stdin.readline

def dijkstra(start:int,N:int,graph:list[tuple[int]])->list[int]:
    INF = float('inf')
    shortcut = [INF]*(N+1)
    shortcut[start] = 0
    ls = [(0,start)]
    while ls:
        length,node = heappop(ls)
        if length > shortcut[node]: continue
        for next_node,added_legth in graph[node]:
            total_length = length + added_legth
            if total_length >= shortcut[next_node]: continue
            shortcut[next_node] = total_length
            heappush(ls,(total_length,next_node))
    return shortcut

# input
N,M = map(int,input().split())
graph = [[] for _ in range(N+1)]
for _ in range(M):
    u,v,w = map(int,input().split())
    graph[u].append((v,w))
    graph[v].append((u,w))
X,Z = map(int,input().split())
P = int(input())
point = list(map(int,input().split()))

# logic
X_shortcut = dijkstra(X,N,graph)
Z_shortcut = dijkstra(Z,N,graph)
ans = float('inf')
for p in point:
    if (X_shortcut[p] == 0) or (Z_shortcut[p] == 0): continue
    ans = min(ans,X_shortcut[p]+Z_shortcut[p])
print(ans) if ans != float('inf') else print(-1)