문제

방향그래프가 주어지면 주어진 시작점에서 다른 모든 정점으로의 최단 경로를 구하는 프로그램을 작성하시오. 단, 모든 간선의 가중치는 10 이하의 자연수이다.

입력

첫째 줄에 정점의 개수 V와 간선의 개수 E가 주어진다. (1≤V≤20,000, 1≤E≤300,000) 모든 정점에는 1부터 V까지 번호가 매겨져 있다고 가정한다. 둘째 줄에는 시작 정점의 번호 K(1≤K≤V)가 주어진다. 셋째 줄부터 E개의 줄에 걸쳐 각 간선을 나타내는 세 개의 정수 (u, v, w)가 순서대로 주어진다. 이는 u에서 v로 가는 가중치 w인 간선이 존재한다는 뜻이다. u와 v는 서로 다르며 w는 10 이하의 자연수이다. 서로 다른 두 정점 사이에 여러 개의 간선이 존재할 수도 있음에 유의한다.

출력

첫째 줄부터 V개의 줄에 걸쳐, i번째 줄에 i번 정점으로의 최단 경로의 경로값을 출력한다. 시작점 자신은 0으로 출력하고, 경로가 존재하지 않는 경우에는 INF를 출력하면 된다.

예제 입력 1

예제 출력 1

0
2
3
7
INF

Solution

[문제 이해]

노드 정보를 그래프에 저장
다익스트라 알고리즘으로 최단거리 탐색

이해한 내용을 바탕으로 코드를 작성했습니다.

[코드]

import sys
import heapq
input = sys.stdin.readline

def solution():
    # 다익스트라
    def dijkstra(start_node, graph):
        distances = [float('inf')] * (V+1)
        distances[start_node] = 0
        queue = []
        heapq.heappush(queue, (0, start_node))

        while queue:
            cur_distance, cur_node = heapq.heappop(queue)
            
            if distances[cur_node] < cur_distance:
                continue
            
            for next_distance, next_node in graph[cur_node]:
                distance = cur_distance + next_distance
                if distance < distances[next_node]:
                    distances[next_node] = distance
                    heapq.heappush(queue, [distance, next_node])

        return distances

    V, E = input().split()
    # V: 노드 개수, E: 간선 개수 -> 노드가 이어진 선
    V, E = int(V), int(E)
    start = int(input())
    graph = [[] for _ in range(V+1)]

    for _ in range(E):
        # u: 출발 노드, v: 도착 노드, w: 거리
        u, v, w = map(int, input().split())
        graph[int(u)].append([int(w), int(v)])

    for i in dijkstra(start, graph)[1:]:
        if i == float('inf'):
            print('INF')
        else:
            print(i)

solution()

[과정]

최초 구현 시에는 그래프를 list가 아닌 dictionary 자료구조를 사용해서 구현했습니다.

하지만 시간 초과라는 결과가 뜨면서 시간 복잡도를 줄이고자 몇 가지를 변경했습니다.

map(int, data) → int(data)
dictionary → list
함수화
comprehention → multiplication

[변경 이전의 코드]

import sys
import heapq
from collections import defaultdict
input = sys.stdin.readline

# 다익스트라
def dijkstra(start_node, graph):
    # 리스트 컴프리헨션을 사용하여 정의
    distances = [float('inf') for _ in range(V+1)]
    distances[start_node] = 0
    queue = []
    heapq.heappush(queue, (0, start_node))

    while queue:
        cur_distance, cur_node = heapq.heappop(queue)
        
        if distances[cur_node] < cur_distance:
            continue
        
        # 자료구조가 dictionary라서 items() 함수를 사용
        for next_node, next_distance in graph[cur_node].items():
            distance = cur_distance + next_distance
            if distance < distances[next_node]:
                distances[next_node] = distance
                heapq.heappush(queue, [distance, next_node])

    return distances

# 정수 변환을 위해 map() 함수 사용
V, E = map(int, input().split())
start = int(input())
# 처음의 자료구조 dictionary 선언
graph = defaultdict(dict)

for _ in range(E):
    # 정수 변환을 위해 map() 함수 사용
    u, v, w = map(int, input().split())
    # dictionary 데이터 추가 형식
    graph[u][v] = w

for i in dijkstra(start, graph)[1:]:
    if i == float('inf'):
        print('INF')
    else:
        print(i)

[여담]

graph 생성 시에 컴프리헨션을 사용한 이유는 *(곱하기)를 사용할 경우 mutable 데이터인 빈 list의 주소값이 복사되기 때문에 데이터 처리가 정상적으로 되지않아서 컴프리헨션을 사용했습니다.

[Python] 리스트 복사 - 보러가기