백준 5719번: 거의 최단 경로

백준 5719번: 거의 최단 경로

풀이 과정

이 문제 때문에 불침번 때도 노트에 손 코딩했다.. 진짜 4시간은 걸린 거 같다. 처음에는 간선을 아예 딕셔너리에서 지워버리는 식으로 코드를 짰는데 시간 초과가 났다. 그래서 간선이 살아있는지 정보를 저장하는 2차원 리스트를 만들었다. 먼저 다익스트라를 돌려서 각 노드까지의 최단 거리를 구하고, 역추적 BFS를 돌린다. 이때 간선의 가중치와 도착 노드까지의 최단 거리의 합이 출발 노드까지의 최단 거리와 같다면 그 간선은 최단 경로에 속해있는 것이다. 그러한 간선들을 모두 배제하고 다익스트라를 다시 돌리면 거의 최단 경로가 나오게 된다.

정답 코드

import sys
import heapq
from collections import deque, defaultdict
input = sys.stdin.readline


def dijkstra(X):
    dist = [float('inf')] * (N)
    dist[X] = 0

    Q = list()
    heapq.heappush(Q, (0, X))

    while Q:
        SD, SN = heapq.heappop(Q)

        if dist[SN] < SD:
            continue

        for FN, FD in L[SN]:
            d = SD + FD
            if dist[FN] > d and LL[SN][FN]:
                dist[FN] = d
                heapq.heappush(Q, (d, FN))

    return dist

def bfs(X):
    Q = deque()
    Q.append(X)

    while Q:
        SN = Q.popleft()

        for FN, FD in L_reverse[SN]:
            if FD + short[FN] == short[SN]:
                if LL[FN][SN]:
                    Q.append(FN)
                    LL[FN][SN] = 0


while True:
    N, M = map(int, input().split())
    if not N or not M:
        break
    S, D = map(int, input().split())
    L = defaultdict(list)
    L_reverse = defaultdict(list)
    LL = [[0] * (N) for _ in range(N)]
    for i in range(M):
        U, V, P = map(int, input().split())
        L[U].append((V, P))
        L_reverse[V].append((U, P))
        LL[U][V] = 1

    short = dijkstra(S)
    if short[D] == float('inf'):
        print(-1)
        continue
    bfs(D)
    result = dijkstra(S)[D]
    print(result if result != float('inf') else -1)

반성할 점

메모리에 너무 집착하지 말자.