STUDY/Algorithm

[백준] 1753 최단경로 python

sinawi95 2021. 4. 22. 19:17
728x90

www.acmicpc.net/problem/1753

 

1753번: 최단경로

첫째 줄에 정점의 개수 V와 간선의 개수 E가 주어진다. (1≤V≤20,000, 1≤E≤300,000) 모든 정점에는 1부터 V까지 번호가 매겨져 있다고 가정한다. 둘째 줄에는 시작 정점의 번호 K(1≤K≤V)가 주어진다.

www.acmicpc.net

import sys
from heapq import heappop, heappush
input = sys.stdin.readline
N, E = map(int, input().split())
K = int(input())
link = [list() for _ in range(N + 1)]
for _ in range(E):
    n1, n2, w = map(int, input().split())
    link[n1].append((w, n2))

# 0번 지점에서 N번 지점까지 이동하는데 걸리는 최소한의 거리
visit = set()
dist = [3000000] * (N + 1)
dist[K] = 0
q = [(0, K)]
while q:
    # 가중치가 가장 작은 노드 탐색
    w, node = heappop(q)
    if node in visit: continue
    visit.add(node)
    for idx in range(len(link[node])):
        adj_w, adj_node = link[node][idx]
        if adj_node in visit: continue
        # 방문하지 않은 점의 가중치보다 (현재 점 +가중치)의 값이 작으면 갱신
        if dist[adj_node] > dist[node] + adj_w:
            dist[adj_node] = dist[node] + adj_w
            # 가중치 순서 최소힙 추가
            heappush(q, (dist[adj_node], adj_node))

for i in range(1, N + 1):
    if dist[i] == 3000000:
        print("INF")
    else:
        print(dist[i])

heapq를 사용하여 다익스트라 알고리즘을 사용하였다. 다익스트라 알고리즘은 시작점이 주어지면 해당 노드에서 모든 노드까지의 최소 거리를 파악할수 있다.

우선, 각 값들을 받아오고 모든 노드의 거리를 inf(충분히 큰 수)로 초기화한다. (inf가 소수점연산을 사용한다라고 들어서 적당한 정수를 사용했다)

시작점을 0으로 두고 인접한 노드를 탐색하며 값을 갱신한다. 갱신하는 조건은 방문하지 않은 인접 노드의 값보다 현재 노드의 값에 가중치를 더한 것이 적어야한다.

이때 갱신이 되면 그 다음 노드도 갱신을 해야하므로 heap에 넣어준다. 이때 heap을 사용하는 이유는 항상 최소인 값을 빼올수 있기 때문이다. 

heapq는 기본값이 min heap이고 원소의 값으로 정렬한다. 튜플로 넣는 경우 맨 앞의 값으로 정렬한다.

해당 노드가 방문했는지를 알아보기위해 set을 사용했는데 리스트는 시간복잡도가 O(N)인데 비해 set과 dictionary는 O(1)이어서 사용했다.

 

처음엔 2차원 인접행렬로 만들었는데 메모리 초과가 떠서 인접리스트로 바꾸었더니 해결하였다.

읽는 값이 많으므로 sys.stdin.readline을 사용해서 속도를 꼭 줄이자...