알고리즘 문제 풀이

16 / 16

백준 5719 거의 최단 경로 with Python

알고리즘다익스트라

문제 링크


알고리즘

  • 그래프 이론
  • 그래프 탐색
  • 다익스트라
  • 최단 경로

풀이

다익스트라 알고리즘을 사용하여 최단 경로를 구한 뒤, 경로를 역추적해서 최단 경로를 그래프에서 제거한다.

나는 그래프를 인접 리스트로 표현했기 때문에, 역추적을 하기 위한 그래프(방향이 거꾸로 된 그래프)를 추가로 만들었다.

(다른 사람들의 풀이를 보니, Node가 최대 500개이므로 그래프를 인접 행렬로 표현하면 메모리 초과가 발생하지 않고 역추적을 위한 그래프를 추가로 만들 필요가 없었다)

다익스트라 알고리즘을 사용하면 출발 노드부터 다른 모든 노드까지의 최단 거리를 구할 수 있다. 역추적할 때는 이 정보들을 이용하여 bfs로 최단 경로를 찾아낼 수 있다. bfs를 수행하는 도중 방문한 경로(edge)를 그래프에서 바로 제거해 주면 된다.

마지막으로, 최단 경로가 제거된 그래프에서 다시 다익스트라 알고리즘을 사용하여 최단 거리를 구하면 문제가 해결된다.

정리하면 다음과 같다.

  1. 다익스트라 알고리즘을 사용하여 출발지부터 다른 모든 노드까지 최단 거리 구하기
  2. 도착지부터 출발지까지 bfs를 사용하여 최단 경로 역추적, 이때 방문한 경로를 그래프에서 제거
  3. 최단 경로가 제거된 그래프에서 다익스트라 알고리즘을 사용하여 최단 거리 구하기

구현 세부사항

역추적 과정에서 edge를 그래프에서 빠르게 제거하기 위해 인접 리스트에 set을 사용했다.

graph = [set() for _ in range(N)]

역추적 과정에서 제거한 경로를 다시 방문하지 않기 위해 해당 노드가 graph에 있는지 확인하는 코드를 추가했다.

if costs[pv] + pc == c and (pc, v) in graph[pv]:
    graph[pv].remove((pc, v))
    q.append((c - pc, pv))

전체 코드

import sys
import heapq
from collections import deque

input = sys.stdin.readline

while True:
    N, M = map(int, input().split())
    if (N, M) == (0, 0):
        exit(0)

    start, end = map(int, input().split())

    graph = [set() for _ in range(N)]  # (cost, v)
    backprop = [set() for _ in range(N)]  # 출발 노드와 도착 노드를 바꾼 그래프

    for _ in range(M):
        a, b, c = map(int, input().split())
        graph[a].add((c, b))
        backprop[b].add((c, a))

    # 첫번째 다익스트라
    q = [(0, start)]
    costs = [float("inf")] * N
    costs[start] = 0

    while q:
        c, v = heapq.heappop(q)
        for nc, nv in graph[v]:
            if c + nc < costs[nv]:
                costs[nv] = c + nc
                heapq.heappush(q, (c + nc, nv))

    # 최단 거리 경로 삭제
    q = deque()
    q.append((costs[end], end))

    while q:
        c, v = q.popleft()
        for pc, pv in backprop[v]:
            if costs[pv] + pc == c and (pc, v) in graph[pv]:
                graph[pv].remove((pc, v))
                q.append((c - pc, pv))

    # 두번째 다익스트라
    q = [(0, start)]
    costs = [float("inf")] * N
    costs[start] = 0

    while q:
        c, v = heapq.heappop(q)
        for nc, nv in graph[v]:
            if c + nc < costs[nv]:
                costs[nv] = c + nc
                heapq.heappush(q, (c + nc, nv))

    print(costs[end] if costs[end] != float("inf") else -1)