Dijkstra 알고리즘으로 모든 정점으로의 최단 경로 구하기

최초작성일 [2020.01.19]

문제 - 백준 1753

다익스트라 알고리즘을 사용하여 어느 한 정점에서 모든 정점으로의 최단경로를 구하는 문제다. 한 가지 유의할 점은 서로 다른 두 정점 사이에 여러개의 간선이 존재할 수 있다는 것이다.

풀이 방법

• 출발 정점에서 각 모든 정점으로 향하는 최단경로를 저장하기 위해 dist 배열을 사용한다.

• 출발 정점에서 v로 향하는 최단 경로를 dist[v]라고 했을 때 (dist[v], v) 형식으로 우선순위 큐에 값을 저장한다.

시간초과 원인

큐를 리스트로 구현하고 known region에 있는 정점을 pop할 때 마다 sorted() 함수로 리스트를 정렬했다. 파이썬의 sorted() 는 시간복잡도가 O(nlogn)이다.

시간초과 해결 방안

heapify(q)를 사용하여 리스트를 힙 트리로 만들고, heappop(), heappush()를 이용하여 큐에서 값을 넣고 꺼낸다. 힙 트리로 구현한 우선순위 큐에서 값을 꺼내고 삽입하는 연산의 시간복잡도는 O(logn)이므로 파이썬의 sorted() 보다 효율적이다.

from heapq import heappop, heappush, heapify
from sys import stdin, maxsize


def relax(v, adj_v, w, dist):
    if dist[adj_v] > dist[v] + w:
        dist[adj_v] = dist[v] + w
        return True
    return False


def dijkstra(start, adjacent_list, n):
    dist = [INF for _ in range(n)]
    dist[start] = 0

    q = list()
    heapify(q)
    heappush(q, (dist[start], start))

    while q:                                        # O(|V|)
        dist_v, v = heappop(q)                      # O(logn)
        for adj_v, adj_w in adjacent_list[v]:       # O(|E|/|V|)
            is_relaxed = relax(v, adj_v, adj_w, dist)
            if is_relaxed:
                heappush(q, (dist[adj_v], adj_v))   # O(logn)

    return dist


INF = maxsize

n, m = list(map(int, stdin.readline().split()))
start = int(stdin.readline().strip('\n'))
start -= 1  # 0-index based

adjacent_list = [list() for _ in range(n)]
for _ in range(m):
    v1, v2, w = list(map(int, stdin.readline().split()))
    adjacent_list[v1-1].append((v2 - 1, w))

result = dijkstra(start, adjacent_list, n)

for i in range(n):
    print('INF') if result[i] == INF else print(result[i])

주목할 점

처음에 모든 정점의 (dist[v], v)를 우선순위 큐에 한꺼번에 넣지 않았다. 즉, known region으로 확정된 start 정점만 큐에 먼저 삽입하고, 탐색을 진행하면서 인접한 정점들 중에 relax된 정점만 큐에 삽입한다. 이 방식으로 구현하면 다음과 같은 장점이 있다.

• 한꺼번에 모든 정점에 대한 정보를 큐에 넣어놓으면 relax를 통해 정점의 dist가 변경될 때마다 큐의 우선순위를 바꿔야하는 복잡한 문제가 생긴다. (힙 트리에서 pop/push 없이 우선순위를 바꾸는 연산을 직접 구현해야 함.) 이러한 문제를 방지할 수 있다.

• 두 정점 사이에 여러개의 간선이 있어도 가중치가 가장 작은 간선에 의해 한 번 relax된 정점은 또 다시 relax될 일이 없다.

시간복잡도

• 큐에는 1개의 정점이 최소 1번 이상 push되고 pop되므로 |V|만큼 pop, push, relax를 실행한다.

• 큐에서 정점 1개를 pop하고 push하는 연산의 시간복잡도는 O(log|V|) 이다.

• 정점 1개는 평균적으로 |E|/|V|개의 간선(인접한 정점)을 가지고 있고 인접한 정점을 대상으로 relax를 시도한다.

따라서 위 알고리즘의 시간복잡도는

정점의 갯수 * {정점 1개를 pop/push 하는 연산의 시간복잡도 + 인접한 정점을 대상으로 하는 relax 연산의 시간복잡도}

|V| * { O(log|V|) + O(|E|/|V|) } = O(|V| * log|V| + |E|) 이다.