백준 1753번 Dijkstra

서론

1753번: 최단경로

이제 BFS, DFS 문제를 어느 정도 풀어보고 최단 경로 문제로 왔습니다. 경로(map이나 graph)가 주어지고 탐색한다는 점이 어떻게 보면 BFS, DFS랑 베이스는 같습니다. 하지만 weight가 있다든지, 특정 조건을 만족시켜야 한다든지의 조건 때문에 일반적인 visited 리스트를 활용한 탐색은 아닌 것 같습니다.

이번 문제는 흔히 알려진 Dijkstra 알고리즘이었습니다. 알고리즘 자체는 알고 있었음에도 구현하는데 뭔가 시간이 오래 걸렸고 힙을 사용해 시간복잡도를 줄여야 한다는 점을 처음에는 생각해내지 못했습니다. 그래도 다른 사람의 코드를 보고 짠 건 아니고 힙이 들어간 알고리즘 자체를 다시 숙지하고 구현했더니 성공했습니다.

제한 조건

• 시간제한 - 1초

시간제한은 1초입니다. 일반적인 컴퓨터로 1초에 1억 번 연산한다는 기준을 가지고 풉니다. 처음에 Dijkstra 알고리즘을 구현할 때 첫 번째 노드에서 인접노드 중 가장 작은 가중치를 가진 인접노드를 선택해서 다음 노드를 선택하며 모든 노드를 탐색했는데 이 과정에서 가장 작은 가중치를 가진 노드를 찾을 때 n의 탐색을 하고 모든 노드를 보며 n의 탐색을 했기 때문에 O(n^2)로 시간초과로 실패했습니다.

그래서 고민을 하다가 힙을 사용해야 된다는 것을 알았습니다. 가장 작은 가중치를 가진 인접노드를 선택하기 위해서 힙구조를 사용하면 시간복잡도가 logn으로 줄기 때문에 전체코드는 O(nlogn)의 시간복잡도를 가지고 제한을 통과할 수 있습니다.

• 메모리 제한 - 256MB

메모리 제한은 256MB입니다. 문제를 풀기 위해선 경로들의 정보를 담고 있는 그래프와 시작노드에서의 최단 거리를 담고 있는 리스트가 필요합니다. 그래프는 한 노드를 기준으로 인접노드와 그에 대한 가중치 정보를 담고 있으니 (20,000+300,000*2)*28byte의 크기를 가지고 최단 거리를 담고 있는 리스트는 20,000*28byte의 크기를 가집니다. 총 18MB 정도의 크기를 가지기 때문에 메모리 제한은 문제가 없습니다.

Dijkstra

시작노드를 기준으로 모든 노드까지의 최단 거리를 구해야 하는 문제입니다. 문제를 해결하기 위한 Dijkstra 알고리즘의 기본적인 동작은 다음과 같습니다. 시작 정점에서의 최단 거리를 담고있는 최단거리 리스트를 사용합니다.(dist_list[vertex_num])

1. 큐에서 노드 한개를 꺼낸다.
2. 꺼낸 노드까지의 거리가 현재 저장된 최단 거리보다 크면 다음 노드를 꺼낸다. (그렇지 않을 때까지)
3. 꺼낸 노드까지의 최단 거리 + 인접노드의 가중치를 계산한다.
4. 위의 값이 현재 인접노드가 가지고 있는 최단 거리보다 작으면 최단 거리를 갱신한다.
5. 최단 거리를 갱신한 인접노드 중 최단 거리가 가장 짧은 노드를 다음 노드로 선정한다.

4번에서 노드를 선정하기 위한 탐색을 할 때 힙을 사용해야 합니다. 처음부터 힙에 넣으면 굳이 탐색하지 않아도 맨 위의 값을 꺼내면 가장 짧은 노드가 나옵니다. 전체 코드는 아래와 같습니다.

import sys
from collections import deque as dq
import heapq
input = sys.stdin.readline

vertex_num, edge_num = map(int, input().split())
vertex_start = int(input())
dist_list = [sys.maxsize]*(vertex_num+1)
dist_list[vertex_start] = 0


def make_graph(vertex_num, edge_num):
    graph = {}
    for i in range(1, vertex_num+1):
        graph[i] = []
    for i in range(edge_num):
        start, end, weight = map(int, input().split())
        graph[start].append([end, weight])
    return graph


def Dijkstra(graph, dist_list):
    heap = []
    heapq.heappush(heap, [0, vertex_start])

    while heap:
        base_dist, base_node = heapq.heappop(heap)
        if base_dist > dist_list[base_node]:
            continue
        for next_node, weight in graph[base_node]:
            next_dist = dist_list[base_node] + weight
            if dist_list[next_node] > next_dist:
                dist_list[next_node] = next_dist
                heapq.heappush(heap, [next_dist, next_node])

    return dist_list[1:]


graph = make_graph(vertex_num, edge_num)
dist_list = (Dijkstra(graph, dist_list))

for dist in dist_list:
    print('INF' if dist == sys.maxsize else dist)

생각

우선은 메모리 제한에 대한 조건을 위와 같은 방법으로 생각하고 있는데, 가끔 문제를 풀다 보니까 이렇게 초깃값에 대한 메모리 말고 알고리즘 자체가 틀려서 메모리 초과가 나는 경우도 있었습니다. 이 문제에는 heap에 데이터가 계속 쌓여서 메모리 초과가 났었고, 다른 경험으로는 함수가 계속 콜 되어 stack이 계속 쌓여서 메모리 초과가 난 적도 있었습니다. 초깃값을 담을 메모리를 설계하는 것도 중요하지만 결론적으론 문제에 대한 해답 자체를 잘 세워서 메모리 초과가 안 나게 하는 것도 중요할 것 같습니다.