[PS] BOJ 5719 / 거의 최단 경로

[PS] BOJ 5719 / 거의 최단 경로
문제 링크: https://www.acmicpc.net/problem/5719
Thumbnail: Photo by D koi (Unsplash)

다익스트라를 활용하는 문제인데, 아이디어의 구현이 어려웠습니다. :

풀이

아이디어

  1. 다익스트라로 시작점 S에서 다른 모든 정점까지의 최단 거리를 구합니다.
  2. 1에서 구한 최단 경로 중, S -> D의 최단 경로에 포함되는 간선들을 삭제합니다.
  3. 다시 다익스트라로 시작점 S에서 D로 가는 최단 경로를 계산해 출력합니다. 새로 계산된 경로가 '거의 최단 경로'이며, 만약 D에 도달할 수 없다면 -1을 출력합니다.

다익스트라

다익스트라 알고리즘은 그래프의 한 정점에서 다른 모든 정점으로의 최단 거리를 구하는 알고리즘입니다. 자세한 구현은 이전 글을 참고해주세요 😉

최단 경로에 포함되는 간선 찾기

다익스트라 알고리즘은 단순히 임의의 정점에서 다른 모든 정점의 최단 거리만 계산합니다. 이 문제의 풀이를 위해선 각 정점을 최단 거리로 방문하기 위해 어떤 간선을 거치는지 알아야 하는데, 이를 어떻게 찾아낼 수 있을까요?

다익스트라 알고리즘의 구현을살펴봅시다. 시작 정점부터 시작해, 각 간선을 통해서 다른 정점과의 최단 거리를 갱신합니다. 그렇다면, 간선을 방문할 때 기록해주면 되지 않을까요?

DIST_MAX = float("inf")
edges = [{} for _ in range(500)]
distance = [DIST_MAX] * 500

def dijkstra(start):
    distance[start] = 0
    pq = []
    heappush(pq, (0, start))

    while pq:
        cost, u = heappop(pq)
        if distance[u] < cost:
            continue
        for v in edges[u]:
            new_dist = cost + edges[u][v]
            if new_dist < distance[v]:
                distance[v] = new_dist
                heappush(pq, (new_dist, v))

다익스트라 알고리즘의 구현

이를 위해, path라는 배열을 만들고 각 정점에 최단 거리로 도달하기 위해 지나간 직전 정점을 기록했습니다.

  • 임의의 정점 $v$로 가는 더 짧은 거리를 찾았다면, path[v]를 빈 배열로 초기화하고 이전 정점 $u$를 기록합니다.
    이전에 기록했던 정점들이 더 이상 최단 거리가 아니기 때문입니다.
  • 임의의 정점 $v$로 가는 기존의 최단 거리와 동일한 거리를 또 찾은 경우에는 path[v]에 이전 정점 $u$를 추가로 기록해줍니다. 한 정점에 도달할 수 있는 최단 경로는 여러 개 존재할 수 있기 때문입니다.
def dijkstra(start):
    distance[start] = 0
    pq = []
    heappush(pq, (0, start))

    while pq:
        cost, u = heappop(pq)
        if distance[u] < cost:
            continue
        for v in edges[u]:
            new_dist = cost + edges[u][v]
            if new_dist < distance[v]:
                distance[v] = new_dist
                path[v].clear()
                path[v].append(u)
                heappush(pq, (new_dist, v))
            elif new_dist == distance[v]:
                path[v].append(u)

다익스트라로 최단 경로 기록하기

최단 경로 역추적하기

이제, path 배열에 기록된 정보를 토대로 최단 경로에 포함된 간선들을 지워야 합니다.

path[v]는 $v$에 도달하기 전에 방문한 정점들의 배열입니다. 다시 말해, 최단 경로를 역방향으로 기록한 배열로 생각할 수 있습니다. 도착점 $D$에서 출발해, $S$까지 path를 따라가며 방문한 정점들을 읽어내면 최단 경로를 찾아낼 수 있습니다.

이 과정은 그래프 탐색으로 진행하면 됩니다. (path를 또 다른 그래프로 생각하기)
BFS, DFS 어느 쪽이든 상관없습니다. 저는 DFS로 구현했습니다.

정점의 방문 여부는 2차원 배열로 기록했습니다. 앞서 한 개의 정점에 도달할 수 있는 최단 경로가 둘 이상일 수 있다고 설명했습니다. 역추적으로 최단 경로를 따라갈 때에도 같은 정점을 두 번 이상 방문할 수 있습니다.

그래프 탐색에서 방문 여부를 기록하는 것은 같은 경우를 반복해서 처리하지 않기 위함인데, 같은 정점에 다른 경로(간선)로 도달하는 것은 중복이 아니므로 모두 처리해야 합니다. 따라서, 방문 여부를 기록하는 배열 check를 2차원 배열로 정의하고, check[출발 정점][도착 정점]으로 방문 여부를 기록했습니다.

또한, 그래프의 간선을 편리하게 지우기 위해서 각 정점에서 출발하는 간선의 정보는 Map에 기록했습니다.

check = [[False] * 500 for _ in range(500)]

def dfs(v):
    if v == S:
        return

    for u in path[v]:
        if not check[u][v]:
            try:
                del edges[u][v]
            except KeyError:
                pass
            check[u][v] = True
            dfs(u)

최단 경로에 포함된 간선 지우기

전체 코드

from heapq import heappop, heappush
input = open(0).readline

DIST_MAX = float("inf")
edges = [{} for _ in range(500)]
distance = [DIST_MAX] * 500
path = [[] for _ in range(500)]
check = [[False] * 500 for _ in range(500)]

def dijkstra(start):
    distance[start] = 0
    pq = []
    heappush(pq, (0, start))

    while pq:
        cost, u = heappop(pq)
        if distance[u] < cost:
            continue
        for v in edges[u]:
            new_dist = cost + edges[u][v]
            if new_dist < distance[v]:
                distance[v] = new_dist
                path[v].clear()
                path[v].append(u)
                heappush(pq, (new_dist, v))
            elif new_dist == distance[v]:
                path[v].append(u)

def dfs(v):
    if v == S:
        return

    for u in path[v]:
        if not check[u][v]:
            try:
                del edges[u][v]
            except KeyError:
                pass
            check[u][v] = True
            dfs(u)

while True:
    N, M = map(int, input().split())
    if N == 0: # Test case ends.
        break
    S, D = map(int, input().split())

    # reset variables
    for i in range(N):
        for j in range(N):
            check[i][j] = False
        distance[i] = DIST_MAX
        path[i].clear()
        edges[i].clear()
    
    for _ in range(M):
        u, v, p = map(int, input().split())
        edges[u][v] = p
    
    # First dijkstra to find the shortest path
    dijkstra(S)
    
    # Remove the shortest path edges
    dfs(D)

    # Second dijkstra to find the almost shortest path
    for i in range(N):
        distance[i] = DIST_MAX
    dijkstra(S)

    print(distance[D] if distance[D] < DIST_MAX else -1)

solution.py