[PS] BOJ 1504 / 특정한 최단 경로

[PS] BOJ 1504 / 특정한 최단 경로
문제 링크: https://www.acmicpc.net/problem/1504
Thumbnail: Photo by Toa Heftiba (Unsplash)

문제 요구 조건을 생각해보면, 결국 최단 경로 문제입니다.

풀이

다익스트라(Dijkstra)!

정점 \(1\)부터 정점 \(N\)까지의 경로 중 임의의 두 정점 \(v_1, v_2\)를 거치는 최단 거리의 경로의 거리(가중치)를 구하는 문제입니다.

먼저 입력을 받아 그래프를 만들어야 합니다. 이번 문제에서는 무방향 그래프라고 명시했으므로, 두 정점 uv에 대해 모두 간선을 추가해주면 됩니다.

N, E = map(int, input().split())
edges = [[] for _ in range(N + 1)]

for _ in range(E):
    u, v, w = map(int, input().split())
    edges[u].append((v, w))
    edges[v].append((u, w))

인접 리스트로 구현한 그래프

이제 다익스트라 알고리즘을 구현해야 합니다. 우선순위 큐 기반의 다익스트라 알고리즘을 구현했습니다. 상세한 설명은 이전 문제의 풀이를 참고해 주세요 😉

INF = 200_000_000 # 방문 불가능한 정점의 거리를 나타낼 상수

def dijkstra(start):
    distance = [INF for _ in range(N + 1)] # 출발 정점 S로부터 각 정점 사이의 최단거리를 저장할 배열
    distance[start] = 0                   # 시작 지점에서 시작 지점까지의 거리는 0이다.

    pq = []                           # 우선순위 큐로 사용할 배열 선언
    heappush(pq, (0, start))              # (거리, 정점)
    while pq:
        dist, node = heappop(pq)      # 우선순위 큐에서 원소 꺼내오기
        if distance[node] < dist: 
            continue
        for next_node, weight in edges[node]: # 다음 노드를 순회하며 우선순위 큐에 넣기
            new_dist = dist + weight
            if new_dist < distance[next_node]:
                distance[next_node] = new_dist
                heappush(pq, (new_dist, next_node))
    return distance

다익스트라 알고리즘

이후 두 정점 \(v_1, v_2\)를 지나는 최단 경로를 계산해야 합니다. 아래 코드에서는 편의상 A, B의 변수에 두 정점을 저장했습니다.

정점 \(1\)에서 정점 \(N\)으로 가는 경로 중 \(v_1, v_2\)를 거치는 최단 경로는 결국 아래 경우 중 거리(가중치)의 합이 최소가 되는 경우입니다.

  • \(1\) -> \(v_1\) -> \(v_2\) -> \(N\)
  • \(1\) -> \(v_2\) -> \(v_1\) -> \(N\)

이를 계산하기 위해, 다익스트라 알고리즘을 총 3번 수행했습니다.

  • 정점 \(1\)에서 다른 모든 정점들로의 거리
  • 정점 \(N\)에서 다른 모든 정점들로의 거리
  • 정점 \(v_1\)에서 다른 모든 정점들로의 거리
    • 문제에서 주어진 그래프가 무방향 그래프이기 때문에 \(v_1\)에서 \(v_2\)로 가는 최단 거리는 \(v_2\)에서 \(v_1\)로 가는 최단 거리와 같습니다. 따라서, 굳이 두 정점 모두에 대해 다익스트라 알고리즘을 수행할 필요는 없습니다.
A, B = map(int, input().split())
dist_start = dijkstra(1)
dist_A = dijkstra(A)
dist_end = dijkstra(N)

dist = min(dist_start[A] + dist_A[B] + dist_end[B], dist_start[B] + dist_A[B] + dist_end[A])
print(-1 if dist >= INF else dist)

결과 출력

전체 코드

from heapq import heappush, heappop
INF = 200_000_000 # 방문 불가능한 정점의 거리를 나타낼 상수
input = open(0).readline
N, E = map(int, input().split())
edges = [[] for _ in range(N + 1)]

# 간선이 없다면 언제나 최단 경로는 존재하지 않는다.
if E == 0:
    print(-1)
    exit()

for _ in range(E):
    u, v, w = map(int, input().split())
    edges[u].append((v, w))
    edges[v].append((u, w))

def dijkstra(start):          
    distance = [INF for _ in range(N + 1)] # 출발 정점 S로부터 각 정점 사이의 최단거리를 저장할 배열
    distance[start] = 0                   # 시작 지점에서 시작 지점까지의 거리는 0이다.

    pq = []                           # 우선순위 큐로 사용할 배열 선언
    heappush(pq, (0, start))              # (거리, 정점)
    while pq:
        dist, node = heappop(pq)      # 우선순위 큐에서 원소 꺼내오기
        if distance[node] < dist: 
            continue
        for next_node, weight in edges[node]: # 다음 노드를 순회하며 우선순위 큐에 넣기
            new_dist = dist + weight
            if new_dist < distance[next_node]:
                distance[next_node] = new_dist
                heappush(pq, (new_dist, next_node))
    return distance

A, B = map(int, input().split())
dist_start = dijkstra(1)
dist_A = dijkstra(A)
dist_end = dijkstra(N)

dist = min(dist_start[A] + dist_A[B] + dist_end[B], dist_start[B] + dist_A[B] + dist_end[A])
print(-1 if dist >= INF else dist)

solution.py