백준 1504 : 특정한 최단 경로

최단 경로라는 문제의 조건과 주어지는 간선들의 비용이 모두 양수이므로, 문제를 잠깐만 살펴보더라도, 다익스트라 문제임을 알 수 있었다. 다익스트라 알고리즘을 손에 익을 만큼 써보면서 연습하고 있기에 이 문제를 풀면서 다시한번 더 다익스트라 알고리즘을 정리해보고자 한다.

<다익스트라 알고리즘>

양의 값을 가진 간선들에 한해서 한 노드에서 다른 모든 노드들까지의 최단거리를 구하는데 사용하는 알고리즘이다.

알고리즘에 사용되는 자료구조 및 자료형

graph

그래프에는 해당 인덱스의 노드에 (연결된 노드 번호, 해당 간선의 cost)가 들어간다.

distance

distance 배열에는 시작노드를 기준으로 거리 0, 그외의 배열들에는 inf값으로 초기화 되어 있으며, 그래프를 탐색하며 최단거리값을 갱신한다.

heapq

한 노드를 방문하게 됐을 때, 해당 노드를 기준으로 연결된 모든 노드들까지의 cost를 고려해서 해당 노드를 거쳐서 방문하는게 더 최단거리인 경우 갱신하기 위해 사용된다.
이때, 힙 자료구조를 사용하면 현재 노드를 기준으로 가장 거리가 가까운 노드부터 탐색할 수 있기 때문에, 가장 가까운 거리를 탐색하는 과정을 O(n)에서 O(log N)으로 단축시킬 수 있다.

전체적인 메커니즘은 다음과 같다.

• 1) 그래프, 시작노드를 인풋으로 받음

• 2) 전체 distance matrix를 시작노드만 0, 나머지는 inf로 초기화

• 3) 거리, 시작노드를 heap에 push

heapq.heappush(queue, (0, start))

• 4) queue가 빌때까지(최단 거리를 고려하며 탐색할 수 있는 모든 노드들은 다 고려할때까지) 탐색

• 5) queue에서 가장 먼저 들어온 값을 pop, 이때 큐에서 pop되는 값은 노드, 거리인 그래프와는 달리, "거리, 해당 노드"임

  why?
  heap 자료구조를 사용하기 위해
  현재 노드에서 가장 가까운 노드가 힙에서 pop되었음을 알 수 있음.

dist, now = heapq.heappop(queue)

• 6) 현재 노드를 거쳐서 방문하는걸 고려하지 않아도 되는 경우 즉, 시작노드에서 지금 다루고 있는 노드까지의 거리가 추가로 고려해야하는 사항들이 남아 있는 dist보다도 짧은 경우에는 더이상 탐색하지 않음

if distance[now] < dist:
	continue

• 7) 현재 노드와 연결된 모든 노드들을 고려해서, 현재노드를 거친 후 해당 노드까지 가는게 더 최단 거리인 경우에는 distance 배열 갱신 후 heap에 push

for i in graph[i]:
#이때 i는 노드, 거리
	cost = dist + i[1]
    if distance[i[0]] > cost:
		distance[i[0]] = cost
        heapq.heappush(queue, (cost, i[0]))

문제 풀이

이렇게 다익스트라 알고리즘을 정리했으니 본격적으로 응용문제를 풀어보도록 하자.
본 문제는 v1, v2의 노드를 반드시 이용해야 한다는 조건이 추가되어 있다. 따라서 1에서 n노드까지의 거리를 구하는 과정 중간에 v1, v2를 반드시 거쳐야 하며, 이를 통해 총 2가지의 경우의 수가 존재함을 알 수 있다.

경우의 수 1
시작노드에서 -> v1 -> v2 -> N으로 가는 경우
시작노드에서 -> v2 -> v1 -> N으로 가는 경우

따라서 시작노드, v1, v2 각각에 대해 다익스트라 알고리즘을 사용해 해당 노드에서 다른 모든 노드들까지의 최단거리를 구한 후,

one = dijkstra(1)
v1_d = dijkstra(v1)
v2_d = dijkstra(v2)

v1_path = one[v1] + v1_d[v2] + v2_d[n]
v2_path = one[v2] + v2_d[v1] + v1_d[n]

두가지의 경우의 수 중에서 더 짧은 값이 v1,v2를 경유하는 최단거리임을 판단하는 것이 본 문제의 핵심이라고 할 수 있다.

전체 소스코드

import sys, heapq

n, e = map(int, sys.stdin.readline().split())
graph = [[]for i in range(n+1)]

for _ in range(e):
    a, b, c = map(int, sys.stdin.readline().split())
    graph[a].append((b,c))
    graph[b].append((a,c))

v1, v2 = map(int, sys.stdin.readline().split())

inf = int(1e9)

def dijkstra(start):
    queue = []
    heapq.heappush(queue, (0, start))
    distance = [inf] * (n + 1)
    distance[start] = 0

    while queue:
        dist, now = heapq.heappop(queue)

        if distance[now] < dist:
            continue

        for i in graph[now]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(queue, (cost, i[0]))

    return distance

one = dijkstra(1)
v1_d = dijkstra(v1)
v2_d = dijkstra(v2)

v1_path = one[v1] + v1_d[v2] + v2_d[n]
v2_path = one[v2] + v2_d[v1] + v1_d[n]

result = min(v1_path, v2_path)
if result < inf:
    print(result)
else:
    print(-1)