1025-TIL(Dijkstra wt heapq)

Dijkstra

for문으로 구현

힙큐를 사용하지 않는 방법. 구현이 쉽지만 느린 O(V^2) 방법.

INF = int(1e9) #무한을 의미하는 값으로 10억을 설정
 
#노드의 개수, 간선의 개수를 입력받기
n,m = map(int, input().split())
#시작 노드 번호를 입력 받기
start = int(input())
#각 노드에 연결되어 있는 노드에 대한 정보를 담를 리스트 만들기
graph = [[] for i in range(n+1)]
#방문한 적이 있는지 체크하는 목적의 리스트를 만들기
visited = [False]*(n+1)
#최단 거리 테이블을 모두 무한으로 초기화
distance = [INF]*(n+1)
 
#모든 간선 정보를 입력받기
for _ in range(m):
    a,b,c = map(int, input().split())
    #a노드에서 b노드로 가는 비용이 c라는 의미
    graph[a].append((b,c))
 
#방문하지 않은 노드 중에서 가장 최단거리가 짧은 노드의 번호를 반환
def get_smallest_node():
    min_value=INF
    index = 0 #가장 최단 거리가 짦은 노드(인덱스)
    for i in range(1, n+1):
        if distance[i] < min_value and not visited[i]:
            min_value = distance[i]
            index=i
    return index
 
def dijkstra(start):
    #시작 노드에 대해서 초기화
    distance[start]=0
    visited[start]=True
    for j in graph[start]:
        distance[j[0]]=j[1]
    #시작 노드를 제외한 전체 n-1개의 노드에 대한 반복
    for i in range(n-1):
        #현재 최단 거리가 가장 짧은 노드를 꺼내서 방문 처리
        now = get_smallest_node()
        visited[now] = True
        #현재 노드와 연결된 다른 노드를 확인
        for j in graph[now]:
            cost = distance[now]+ j[1]
            #현재 노드를 거쳐서 다른 노드로 이동하는 거리가 더 짧은 경우
            if cost < distance[j[0]]:
                distance[j[0]]=cost
 
#다익스트라 알고리즘 수행
dijkstra(start)
 
#모든 노드로 가기 위한 최단 거리를 출력
for i in range(1, n+1):
    #도달할 수 없는 경우, 무한(INFINITY)이라고 출력
    if distance[i]==INF:
        print("INFINITY")
    #도달할 수 있는 경우 거리를 출력
    else:
        print(distance[i])

이것이 코딩 테스트다 책의 코드인데 이해가 어려워서(e.g. distance[j[0]]=j[1]) 아래 블로그에서 찾은 코드

출처

import sys
input = sys.stdin.readline
INF = int(1e9)

#n, m = 노드 개수, 간선 개수
n, m = map(int, input().split())
start = int(input()) #시작 노드
# 각 노드에 대해 (연결 노드, 거리) 값을 담기 위한 그래프
graph = [ [] for _ in range(n+1) ]
#다익스트라 알고리즘 수행 시 방문 여부 체크를 위한 리스트
visited = [0] * (n + 1) 
#최단 거리 테이블 
distance = [INF] * (n + 1)

for _ in range(m):
  #a노드에서 b노드로 가는 간선의 거리 c
  a, b, c = map(int, input().split())
  graph[a].append((b, c))

#최단 거리 테이블의 최소값을 얻기 위한 함수
def get_smallest_node():
  idx = 0 
  min_value = INF
  for i in range(1, n + 1):
    if distance[i] < min_value and visited[i] == 0:
      min_value = distance[i]
      idx = i
  return idx

def dijkstra(start):
  #시작 노드에서 시작 노드까지의 거리 0 저장
  distance[start] = 0
  #시작 노드를 포함한 n개의 노드 방문
  for _ in range(n+1):
    #현재 가장 최단 거리가 짧은 노드를 꺼내서 방문 처리
    cur = get_smallest_node()
    visited[cur] = 1
    #현재 방문한 노드와 연결된 노드들에 대해 최단 거리 테이블 갱신
    for next in graph[cur]:
      #현재 노드까지의 최단 거리 + 현재 노드에서 다음 노드 까지의 거리
      cost = distance[cur] + next[1]
      #기존의 다음 노드까지의 최단 거리보다 현재 노드를 거친 다음 노드까지의 거리가 더 짧으면 테이블 갱신 
      if distance[next[0]] > cost:
        distance[next[0]] = cost

dijkstra(start)

for i in range(1, n+1):
  if i == INF:
    print("INFINITY")
  else:
    print(distance[i])

heapq로 구현

힙큐 라이브러리 사용 이유 - 우선순위 큐를 사용하기 위해.

아래는 heapq 라이브러리를 사용해서 get_smallest_node() 를 사용하지 않아도 되는 모습이다. 힙큐를 사용하면 시간복잡도가 O(ElogV)가 보장된다.

왜냐하면 for문으로 구현한 다익스트라 알고리즘은 최단거리가 가장 짧은 노드를 탐색하기 위해서 매번 최단 거리 테이블을 선형적으로 탐색해야 했다. 이 과정에서만 O(V)의 시간이 걸렸다.

때문에 힙큐 라이브러리를 사용해서 힙 자료구조를 이용하면 선형 시간이 아닌 로그 시간이 걸린다.

힙 자료구조는 우선순위 큐를 구현하기 위하여 사용하는 자료구조 중 하나이다
(파이썬에서 우선순위 큐가 필요할 때 쓸 수 있는 모듈 두개 중 하나가 PriorityQueue, 나머지 하나는 heapq).
PriorityQueue보다 heapq가 더 빠르게 동작하기 때문에 힙큐 라이브러리를 가져다 쓰면 된다.

우선순위 큐는 왜 쓰면 좋은데?

우선순위 큐의 구현 방법은 다양한데 기본적으로 이렇게 두가지로 나뉜다.
1. 리스트 (제일 단순)
2. 힙 (Heap)

우선순위 큐 뭔데?

우선순위가 가장 높은 데이터를 가장 먼저 삭제한다.

우선순위 값을 표현할 때 물건 데이터라고 하면 (가치, 물건)으로 묶여서 우선순위 큐에 넣고, 이후에 우선순위큐에서 물건을 꺼내게 되면(popleft) 항상 가치가 높은 물건이 먼저 나오게 된다.(참고로 기본 힙큐 라이브러리에 규현되어있는 힙소트는 최소힙으로 되어있어서 최대힙이 필요하다면 heapq앞에 -를 이용해서 최대 힙 구현하면 된다.)

최소힙은 값이 낮은 데이터가 먼저 삭제, 최대힙은 값이 큰 데이터가 먼저 삭제된다.

#최대힙
import heapq
def heapsort(iterable):
    h = []
    result = []
    #모든 원소를 차례대로 힙에 삽입
    for value in iterable:
        heapq.heappush(h, -value) #value에 -
    #힙에 삽입된 모든 원소를 차례대로 꺼내서 담기
    for i in range(len(h)):
        result.append(-heapq.heappop(h)) #heapq.heappop(h)에 -
 
    return result
 
result = heapsort([1,3,5,7,9,2,4,6,8,0])
print(result)
#출력
#[9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

import heapq
def heapsort(iterable):
    h = []
    result = []
    #모든 원소를 차례대로 힙에 삽입
    for value in iterable:
        heapq.heappush(h, value)
    #힙에 삽입된 모든 원소를 차례대로 꺼내서 담기
    for i in range(len(h)):
        result.append(heapq.heappop(h))
 
    return result
 
result = heapsort([1,3,5,7,9,2,4,6,8,0])
print(result)
#출력 
#[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

우선순위 큐 그래서 어떻게 쓰는건데?

우선순위 큐를 적용해도 다익스트라 알고리즘이 동작하는 기본 원리는 동일!
때문에 최단 거리를 저장하는 1차원 리스트(최단거리 테이블)은 아까와 같이 그대로 이용하면 된다.
즉, 현재 가장 가까운 노드를 저장하기 위한 목적으로만 우선순위 큐를 추가로 이용한다고 보면 되겠다.

파이썬에서는 튜플 (0, 1)을 우선순위 큐에 넣는다.
heapq 라이브러리는 원소로 튜플을 입력받으면 튜플의 첫 번째 원소를 기준으로 우선순위 큐를 구성한다.
따라서 (거리, 노드) 순으로 튜플 데이터를 구성하면 거리순으로 정렬이 되는 것이다.

import heapq
import sys
input = sys.stdin.readline
INF = int(1e9) #무한을 의미하는 값으로 10억
 
#노드의 개수, 간선의 개수를 입력받기
n,m = map(int, input().split())
#시작 노드 번호를 입력받기
start = int(input())
#각 노드에 연결되어 있는 노드에 대한 정보를 담는 리스트 만들기
graph = [[] for i in range(n+1)]
#최단 거리 테이블을 무한으로 초기화
distance = [INF]*(n+1)
 
#모든 간선 정보를 입력받기
for _ in range(m):
    a,b,c = map(int, input().split())
    #a번 노드에서 b번 노드로 가는 비용이 c
    graph[a].append((b,c))
 
def dijkstra(start):
    q=[]
    #시작 노드로 가기 위한 최단 경로는 0으로 설정하여 큐에 삽입
    heapq.heappush(q, (0, start))
    distance[start]=0
    #q가 비어있지 않다면
    while q:
        #가장 최단 거리인 노드에 대한 정보 꺼내기
        dist, now = heapq.heappop(q)
        #현재 노드가 이미 처리됐다면 skip
        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(q, (cost, i[0]))
 
#다익스트라 알고리즘 실행
dijkstra(start)
 
#모든 노드로 가기 위한 최단 거리 출력
for i in range(1, n+1):
    #도달할 수 없는 경우, 무한 출력
    if distance[i]==INF:
        print("INF")
    else:
        print(distance[i])