Post

[백준] 5972 파이썬

문제


  • 백준 5972 (골드 5)

백준 5972

n개의 노드와 m개의 양방향 간선과 그 가중치가 주어질 때, 1번 노드부터 n번 노드까지의 최소 가중치를 구하는 문제였다.
노드와 간선이 최대 50,000개에 시간 제한이 1초여서 시간복잡도를 고려해야하는 문제였다.

image

내 코드


우선 dfs가 떠올라서 생각나는데로 구현해보았다.
짜면서도 시간제한에 걸릴 것 같다는 생각이 들었는데 아니나 다를까 바로 시간 초과가 발생했다.
다음으로는 다익스트라 알고리즘이 떠올라서 해당 방식으로 구현해보았다.
파이썬의 heapq를 활용하여서 코드를 작성하였다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import sys
import heapq
input = sys.stdin.readline

n, m = map(int, input().split())
cows = [[] for _ in range(n)]
for _ in range(m):
    a, b, c = map(int, input().split())
    cows[a - 1].append((b - 1, c))
    cows[b - 1].append((a - 1, c))
dist = [10**9] * n

def dijkstra(st):
    q = []
    heapq.heappush(q, (0, st))
    dist[st] = 0

    while q:
        d, cur = heapq.heappop(q)

        if dist[cur] < d:
            continue

        for i, j in cows[cur]:
            if d + j < dist[i]:
                dist[i] = d + j
                heapq.heappush(q, (d + j, i))

dijkstra(0)
print(dist[n - 1])

image

This post is licensed under CC BY 4.0 by the author.