[백준-파이썬] 1753-최단경로

문제바로가기

완벽히 다엑스트라 구현 문제이다.

import sys
sys.setrecursionlimit(10**6)

v, e = map(int, sys.stdin.readline().split())
k = int(sys.stdin.readline())
min_path = ['inf']*(v+1)
edge = [[False]*(v+1) for _ in range(v+1)]

for _ in range(e):
    u, ev, w = map(int, sys.stdin.readline().split())
    edge[u][ev] = w

def search(node):
    if min_path[node] == 'inf':
        min_path[node] = 0
    next_node = []
    
    for n in range(v+1):
        if edge[node][n]:
            if min_path[n] == 'inf':
                min_path[n] = min_path[node] + edge[node][n]
                next_node.append([edge[node][n], n])
            else:
                min_path[n] = min(min_path[n], min_path[node]+edge[node][n])
    
    if next_node:
        next_node.sort()
        search(next_node[0][1])

min_path[k] = 0
search(k)

for i in range(k, v+1):
    print(min_path[i])

처음에 구현한 코드는 잘 돌아가지만, 인접리스트로 구현했기 때문에 메모리 초과를 피할 수 없었다. 왜냐하면 v가 커질수록 더 많은 메모리가 필요한데 최대 v는 20000이므로 10^8을 수용할 수 있는 메모리가 필요하다. 굉장히 비효율적이다. 결국 우선순위 큐로 다시 코드를 짰다.

import heapq
import sys
sys.setrecursionlimit(10**6)

INF = int(1e9)

V, E = map(int, sys.stdin.readline().split())
K = int(sys.stdin.readline())
min_path = [INF]*(V+1)
edge = [[] for _ in range(V+1)]

for _ in range(E):
    u, v, w = map(int, sys.stdin.readline().split())
    edge[u].append((v, w))

def dijkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    min_path[start] = 0
    
    while q:
        cur_dis, cur_node = heapq.heappop(q)
        if min_path[cur_node] < cur_dis:
            continue
            
        for node, dis in edge[cur_node]:
            if cur_dis + dis < min_path[node]:
                min_path[node] = cur_dis + dis
                heapq.heappush(q, (cur_dis + dis, node))

dijkstra(K)
for i in range(1, V+1):
    print(min_path[i] if min_path[i] != INF else "INF")