백준 1753 최단경로 파이썬 문제풀이 (python 다익스트라)

문제 보러가기

 

 

백준 최단경로 파이썬 문제풀이

최단경로라는 문제는 그래프에 거리 즉 가중치 값이 있다.

그래프가 있고 최단거리와 같이 edge에 가중치가 있는 경우

 

"다익스트라" 알고리즘을 적용한다.

기본적인 내용은 여기서 다루지 않고 바로 적용해보자.

 

다익스트라 알고리즘에서 유의할 점은 node의 edge연결정보를 가지고 

최단거리를 기록할 때 순서가 중요하다.

 

반영할 노드 순서를 잘못해서 최단거리에 반영하면, 최단거리가 아닌 값을 구하게 된다.

 

from collections import deque

N,M = map(int, input().split())

start = int(input())

links = [[] for _ in range(N+1)]
visited = [0 for _ in range(N+1)]
dist = [ 1e9 for _ in range(N+1) ]

for _ in range(M):
    A,B,C = map(int, input().split())
    links[A].append([B,C])

q = deque()
q.append(start)
visited[start] = 1
dist[start] = 0

def shortest_finder():
    mini = 1e9
    idx = 0
    for i in range(1,N+1):
        if dist[i] <= mini:
            idx = i
            mini = dist[i]
    return idx

while q:
    node = q.popleft()
    for nxt, weight in links[node]:
        # 1. 각각 노드 값 업데이트
        # 2. 여러 경로가 있으면 min 비교
        # 3. 시작점으로부터 거리 보고, 짧은 순서대로 탐색.
        dist[nxt] = min(dist[node] + weight, dist[nxt])
        q.append(nxt)
        visited[nxt] = 1

        idx = shortest_finder()

        if idx in q:
            q.remove(idx)
            q.appendleft(idx)

for i in dist[1:]:
    if i == 1e9:
        print("INF")
    else:
        print(i)

이렇게 다익스트라를 이용해서 최단거리 알고리즘을 만들었지만.

 

시간 초과가 났다.

더 빠르게 구하는 방법을 탐색해보자. 

 

우선순위 큐 heapq

파이썬의 heapq라는 자료 구조를 사용하면, 

우선순위 큐를 사용할 수 있다.

 

이를 이용해서 원래 거리가 가장 가까운 노드를 찾았던 로직의

시간을 절약시킬 것이다. 

 

거리에 대해서 우선순위 큐를 이용한 로직으로 변경해보자. 

 

정답 코드

from collections import deque
import heapq
N,M = map(int, input().split())

start = int(input())

links = [[] for _ in range(N+1)]
dist = [ 1e9 for _ in range(N+1) ]

for _ in range(M):
    A,B,C = map(int, input().split())
    links[A].append([B,C])

q = []
q.append([0,start])
heapq.heappush(q, [0,start])
dist[start] = 0

while q:
    # 우선순위 큐를 이용한 거리 정렬

    _w, node = heapq.heappop(q)

    for nxt, weight in links[node]:
        # 1. 각각 노드 값 업데이트
        # 2. 여러 경로가 있으면 min 비교
        # 3. 시작점으로부터 거리 보고, 짧은 순서대로 탐색.
        if dist[node] + weight < dist[nxt]:
            dist[nxt] = dist[node] + weight
            heapq.heappush(q, [dist[nxt], nxt])

for i in dist[1:]:
    if i == 1e9:
        print("INF")
    else:
        print(i)

이렇게 heap 정렬에 기반한 우선순위 큐를 이용하면 

 

시간 부족을 이겨내고 정답을 찾을 수 있다.

 

하지만 온전히 집중해서 풀어내질 못해서 나중에 다시 또 풀어보자.