[백준1753번] 최단경로 / Python3

문제

방향그래프가 주어지면 주어진 시작점에서 다른 모든 정점으로의 최단 경로를 구하는 프로그램을 작성하시오. 단, 모든 간선의 가중치는 10 이하의 자연수이다.

입력

첫째 줄에 정점의 개수 V와 간선의 개수 E가 주어진다. (1≤V≤20,000, 1≤E≤300,000) 모든 정점에는 1부터 V까지 번호가 매겨져 있다고 가정한다. 둘째 줄에는 시작 정점의 번호 K(1≤K≤V)가 주어진다. 셋째 줄부터 E개의 줄에 걸쳐 각 간선을 나타내는 세 개의 정수 (u, v, w)가 순서대로 주어진다. 이는 u에서 v로 가는 가중치 w인 간선이 존재한다는 뜻이다. u와 v는 서로 다르며 w는 10 이하의 자연수이다. 서로 다른 두 정점 사이에 여러 개의 간선이 존재할 수도 있음에 유의한다.

출력

첫째 줄부터 V개의 줄에 걸쳐, i번째 줄에 i번 정점으로의 최단 경로의 경로값을 출력한다. 시작점 자신은 0으로 출력하고, 경로가 존재하지 않는 경우에는 INF를 출력하면 된다.

예제 입력 1

5 6
1
5 1 1
1 2 2
1 3 3
2 3 4
2 4 5
3 4 6

예제 출력 1

0
2
3
7
INF

나의 풀이

다익스트라 알고리즘을 사용하여 푸는 문제로 나는 다익스트라 알고리즘은 이해했지만 애를 먹은 문제이다. 처음에 인접행렬로 구현을 하여 메모리 초과를 경험하고, 인접리스트로 수정을 하여 시간초과를 경험하고, 마지막으로 힙을 사용함으로 문제를 해결할 수 있었다.

우선 이 문제를 풀기위해서는 다익스트라(dijkstra) 알고리즘을 이해해야한다. 처음에 보면 복잡해 보이지만 막상 하나하나 들여다보면 그닥 어렵지 않은 알고리즘이다. 다익스트라 알고리즘은 한 정점에서부터 다른 모든 정점들로 가는 최소의 cost를 구하는 알고리즘이라고 할 수 있다.

위의 그래프가 문제의 예제에서 주어지는 그래프를 그림으로 그린 것이다. 문제는 1번 노드에서 2,3,4,5번으로 가는데 드는 최소의 비용을 구하는 문제이다. 따라서 1번에서 어떤 노드를 거치든 갈 수 없는 5번과 같은 경우에는 INF를 출력해주어야한다. 다익스트라 알고리즘을 사용하면 노드의 개수만큼 반복을 하면서 이웃한 노드들을 모두 확인해 각 반복마다 최소의 값을 확정지으면 된다.

위의 배열이 다익스트라 알고리즘의 과정을 나타낸 그림이다. 과정을 설명하자면 아래와 같다.

1. 처음 1번에서 출발하고 cost는 0으로 1번 index의 값이 0으로 채워지고 1번을 방문.

2. 1번과 이웃한 곳의 비용을 최신화 시킨다. 1번에서 2,3번으로 갈 수 있으니 2,3번 인덱스를 각각의 cost로 저장을 해준다. 4,5번은 1번에서 접근이 불가능하니 INF이다. 이번 확인에서는 2번으로 가는 비용이 제일 작으므로 1->2의 최소비용이 확정이 된다. (확정이 되면 다시는 확인해줄 필요가 없게된다.)

3. 이제 세번째 배열에서 3번 인덱스의 값을 정하는 과정을 생각해보자.
    - 1번에서 3번으로는 바로 갈 수도 있지만, 2번을 거쳐서 갈 수도 있다. 이러한 경우를 생각해 (1 -> 3)과 (1 -> 2 -> 3)의 비용을 비교해서 작은 값을 저장하면 된다. 따라서 3 < 6 으로 1 -> 3인 경우의 비용이 저장된다.
    - 2번을 거쳐 4번으로 가는 비용도 (1 -> 2 -> 4)를 7이 저장이 된다.
    - 이번 반복에서 3번이 가장 작은 비용으로 3번을 방문한다.

4. 3번을 방문하였으니 이제 3번과 이웃한 노드들의 값을 최신화 시킨다.
 - 3번에서 갈 수 있는 노드는 4번 뿐이다. 따라서 앞서 구해진 (1 -> 2 -> 4)와 (1 -> 3 -> 4)를 비교해보고 작은 값을 저장해주면 된다. 따라서 7이 저장이 되고 4번을 방문하게 된다.

5. 5번은 어디에서도 접근을 할 수 없으므로 INF로 고정이 되어 다익스트라 알고리즘이 종료된다.

정말로 위와 같은 배열이 나오는지 코드를 짜고나서 각 반복마다 최소 비용을 저장하는 배열을 출력해 확인해보았다.

메모리 초과 문제 해결

필자는 그래프를 저장하기 위해 이해가 쉽고 구현에 쉬운 인접행렬을 사용하였다. 인접행렬은 노드의 수 * 노드의 수 크기의 2차원 배열을 사용하여 간선들의 정보를 저장하는 방법으로 메모리를 많이 잡아먹는다는 단점이 있었다. 따라서 이 문제에서는 인접리스트를 사용하여야 했었다. 인접리스트는 각 노드에 대한 간선들을 해당 인덱스에 저장을 하는 것으로 2차원 배열에 비해서 메모리를 절약할 수 있었다.

시간 초과 문제 해결

처음 구현을 했을때에는 한 번의 반복마다 최소 비용을 저장하는 배열에서의 최솟값을 반복문을 통해서 찾았다. 이러한 경우에 시간복잡도는 O(n)이다. 그러나 우선순위 큐(최소힙)을 사용하면 시간복잡도가 O(logn)으로 기존의 방법보다 작았다. 따라서 최소힙에 [cost, 노드의 번호]의 형태로 넣고 빼줌으로써 최소 비용을 바로 뽑아낼 수 있었다.

코드

# 1753번 최단경로
import sys, heapq

# 최대값(INF) 변수 선언
inf = float('inf')

# dijkstra algorithm
def dijkstra(v,k,g):

    # 최소 거리를 저장하는 배열
    d = [inf for _ in range(v+1)]
    d[k] = 0

    hq = []
    heapq.heappush(hq, [0,k])   # heapq에 [cost,node] 넣음.

    print(d)
    while hq:
        tmp = heapq.heappop(hq)     # cost가 최소인 노드 pop
        cost = tmp[0]
        node = tmp[1]

        if cost > d[node]:
            continue

        # 이웃한 노드들을 확인하여 최솟값 업데이트
        for n in g[node]:
            neighbor = n[0]
            n_cost = d[node] + n[1]

            # node를 지나서 이웃으로 가는 cost가 최소인 경우 최신화
            if n_cost < d[neighbor]:
                d[neighbor] = n_cost
                heapq.heappush(hq, [n_cost,neighbor])
        print(d)

    return d

# main
V, E = [int(x) for x in sys.stdin.readline().split()]

# 시작 노드
K = int(input())

# 인접리스트
g = [[] for _ in range(V+1)]
for _ in range(E):
    v, e, w = [int(x) for x in sys.stdin.readline().split()]
    g[v].append([e,w])

answer = dijkstra(V,K,g)

for dis in answer[1:]:
    print(dis if dis != inf else 'INF')