[백준 1753번] 최단경로 (Python 풀이)

https://www.acmicpc.net/problem/1753

문제

문제 이해

먼저, 문제의 제목처럼 최단경로를 구하는 문제이다. 선뜻, BFS (너비우선탐색)이 생각나지만, 이 문제는 가중치가 있다.

1, 0 으로 이루어진 가중치라면 BFS풀이가 가능하나, 1을 초과하는 가중치가 있기 때문에 이 문제는 BFS로 풀지 못한다.

즉, Dijkstra (다익스트라) 알고리즘을 사용한다.

첫째 줄에는 노드 수, 간선 수가 나온다.

여기서 간선이라고 하는 것은 가중치가 있고 방향이 있는 화살표이다. (양쪽을 향하는 화살표가 아니란 말이다.)

예제 입력을 그림으로 나타내면 저런 형태를 가진다.

 

풀이 방법

1. 먼저 간선(v)과 가중치(e)가 들어가는 edge 배열을 만든다.위의 예제 입력의 edge 배열은edge = [ [ ], [(2, 2), (3, 4)], [(3,4), (4,5)], [(3, 4)], [ ], [ ] ] 

이렇게 된다. 아래 코드처럼 넣어주면 되는 것이다.

edge = [[] for _ in range(V+1)]
for _ in range(E):
    u, v, e = map(int, input().split())
    edge[u].append((v,e))

2. 가중치(거리) 계산을 위해 distance배열을 만든다.

distance = [float('inf')]*(V+1)

무한으로 초기화를 해준다.

 

3. 위와 같은 재료준비가 끝났으니, dijkstra 알고리즘 함수를 작성한다.          

def dijkstra(start):	
    # heapq는 요소의 첫번째의 첫번째만 비교하여 최소값을 찾기 때문에
    # q에 들어가는 순서는 ( 거리(가중치), 노드번호 ) 이다.
    q = [(0, start)]	
    distance[start] = 0	#시작 노드는 자기 자신과 거리가 0이기 때문에 0으로 설정
    while q:
        dis, now = heapq.heappop(q)	# q 에서 현재 노드의 거리와 번호를 뽑아낸다.
        for node, weight in edge[now]:	# edge에서 node와 weight를 뽑는다.
            # 현재 노드의 distance와 연결된 노드의 가중치를 더한 것 (=d) 과 연결된 노드의 distance를 비교
            d = dis + weight			
            if d < distance[node]: # d가 더 작다면 연결된 노드의 distance 는 d가 된다.
                distance[node] = d
                heapq.heappush(q, (d, node))	# q라는 heap에 넣어준다.

4. 출력을 위한 코드를 짜는 것으로 마무리한다.

for i in range(1, V+1):
    if distance[i] == float('inf'):
        print('INF')
    else:
        print(distance[i])

 

전체코드 (메모리 68620 KB, 시간 782 ms)

import heapq
import sys
input = sys.stdin.readline

def dijkstra(start):
    q = [(0, start)]
    distance[start] = 0
    while q:
        dis, now = heapq.heappop(q)
        for node, weight in edge[now]:
            d = dis + weight
            if d < distance[node]:
                distance[node] = d
                heapq.heappush(q, (d, node))

  
V, E = map(int, input().split())
start = int(input())
edge = [[] for _ in range(V+1)]
distance = [float('inf')]*(V+1)
for _ in range(E):
    u, v, e = map(int, input().split())
    edge[u].append((v,e))
dijkstra(start)
for i in range(1, V+1):
    if distance[i] == float('inf'):
        print('INF')
    else:
        print(distance[i])

풀이 후 생각 정리

처음에는 별 생각 없이 deque로 배열을 짜고 sort로 정렬을 하는 것을 사용했는데, 시간 초과가 나왔다.

heap을 사용하는 것이 시간적으로 훨씬 우세하기 때문에, heap을 사용하는 문제였다.

처음 잘못짠 코드는 아래와 같다 (시간초과)

def dijkstra(start):
    q = deque([start])
    distance[start] = 0
    while q:
        c = q.popleft()
        edge[c].sort(key = lambda x:x[1])
        for node, weight in edge[c]:
            d = distance[c] + weight
            if d < distance[node]:
                distance[node] = d
                q.append(node)