Baekjoon 5719. 거의 최단 경로 / Python

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

문제 정리

길의 최단 경로를 제외한 경로들 중의 최단 경로를 찾는다.

문제 풀이 방식

최단거리를 찾고 -> 최단 경로를 제거하고 -> 최단 경로를 찾는다.

라고 하면 간단해 보이지만 생각보다 시간초과가 많이났던 코드

다익스트라 알고리즘을 두 번 사용한다.

import sys
from heapq import heappop, heappush
from collections import deque as dq
input = sys.stdin.readline 

inf = 1e10 

def dijkstra():
    dp = [inf] * n  # 최단 거리 저장할 배열
    q = []
    heappush(q, (0, s, s))  # 시작점의 거리 0으로 설정하고 큐에 추가
    while q:
        dist, start, end = heappop(q)  # 현재 거리, 현재 노드, 이전 노드
        if dp[start] < dist:
            continue  # 이미 처리된 노드면 스킵
        if dp[start] == dist:
            path[start].append(end)  # 동일한 최단 거리 경로 추가
            continue
        dp[start] = dist
        path[start] = [end]  # 최단 거리 갱신 및 경로 저장
        for i in range(n):
            if arr[start][i] != inf:  # 연결된 모든 도로에 대해
                heappush(q, (dist + arr[start][i], i, start))  # 새로운 거리 계산 후 큐에 추가
    return dp[d]  # 도착점까지의 최단 거리

def destroy():
    q = dq([(d, d)])  # 도착점에서 시작
    while q:
        x, y = q.popleft()  # 현재 노드, 이전 노드
        if arr[x][y] == inf:
            continue  # 이미 제거된 도로면 스킵
        arr[x][y] = inf  # 도로 제거
        for k in path[x]:  # 현재 노드로부터의 경로 추적
            q.append((k, x))  # 이전 노드와 연결된 노드 큐에 추가

while True:
    n, m = map(int, input().split())  # 장소 수와 도로 수 입력
    res = inf  # 초기 결과값을 무한대로 설정
    if not n + m:
        break  # 입력의 마지막 줄에 도달하면 종료
    s, d = map(int, input().split())  # 시작점과 도착점 입력
    arr = [[inf] * n for _ in range(n)]  # 인접 행렬 초기화
    arr[d][d] = 0  # 도착점 자기 자신으로 가는 경로는 0으로 설정

    for _ in range(m):
        u, v, p = map(int, input().split())  # 도로 정보 입력
        arr[u][v] = p  # 도로 길이 설정
    path = {}  # 최단 경로 정보 저장할 딕셔너리 초기화
    short = dijkstra()  # 다익스트라 알고리즘을 사용하여 최단 경로 찾기
    if short == inf:
        print(-1)  # 최단 경로가 없으면 -1 출력
    else:
        destroy()  # 최단 경로 제거
        short = dijkstra()  # 다시 다익스트라 알고리즘을 사용하여 거의 최단 경로 찾기
        if short == inf:
            print(-1)  # 거의 최단 경로가 없으면 -1 출력
        else:
            print(short)  # 거의 최단 경로의 길이 출력