[백준 / 1753] 최단경로 - Python

시간 초과와 여러개의 간선에서 문제가 있었다.

 

최단 경로는 '다익스트라(Dijkstra) 알고리즘'을 먼저 공부해야 한다.

↳ 출발 지점으로부터 비용(가중치)이 가장 낮은 경로를 계산

 

2021.03.18 - [# Coding/# Algorithm] - [보기 쉬운 Algorithm] 다익스트라(Dijkstra) 알고리즘 - Python

[보기 쉬운 Algorithm] 다익스트라(Dijkstra) 알고리즘 - Python

<다익스트라 알고리즘>

 

import sys
import heapq

def dijkstra(graph, start):
    distances = {node: float('INF') for node in graph}
    distances[start] = 0
    queue = []
    heapq.heappush(queue, [distances[start], start])

    while queue:
        current_distance, node = heapq.heappop(queue)

        if distances[node] < current_distance:
            continue

        for new_node, new_distance in graph[node].items():
            distance = new_distance+current_distance
            if distance < distances[new_node]:
                distances[new_node] = distance
                heapq.heappush(queue, [distance, new_node])

    return distances

 

입력 시간과 알고리즘 시간을 줄이기 위해서 sys와 heqpq를 import한다.

distances - {노드 : 최단거리} 

distances[start]는 계산 되면 안되므로 0으로 만들어 준다.

queue에는 현재 기준 노드와 그 노드까지 온 거리가 들어간다.

 

queue에 값이 없을 때 까지 반복한다.

현재 기준 노드와 그 현재 노드까지의 비용(가중치)를 받아온다.

distances에 있는 거리보다 길면 무시한다.

 

짧을 경우 graph에 있는 현재 노드에서 갈 수 있는 노드와 비용(가중치)을 받아온다.

다음 노드의 길이를 구하기 위해서 현재 노드 길이와 앞으로의 노드 길이를 더해준다.

그 길이가 distances에 있는 값보다 짧을 경우 값을 바꿔준다.

그리고 queue에 다음 노드와 다음 노드까지 가는데 걸린 길이를 넣어준다.

 

distances에 모든 노드의 최소 값이 담겨 있으므로 반환해준다.

 

<입력1>

 

V, E = map(int, sys.stdin.readline().split())
K = int(sys.stdin.readline())

 

V(node 개수), E(간선 개수), K(출발점) 변수를 입력받는다.

<그래프 선언>

 

graph = {}
for i in range(1, V+1):
    graph[i] = {}

 

간선이 없는 노드도 있을 수 있으니, 모든 노드를 그래프에 추가해준다.

 

<입력2>

for _ in range(E):
    u, v, w = map(int, sys.stdin.readline().split())
    
    if v in graph[u]:
        if graph[u][v] > w:
            graph[u][v] = w
    else:
        graph[u][v] = w

 

간선의 개수(E)만큼 입력을 받는다.

문제에 노드 사이에 여러 간선이 있을 수 있으니 중복 간선에서 가장 비용(가중치)가 낮은 값만 저장한다.

if문에서는 그래프에 u노드에 v로 가는 간선이 없으면 새로 추가해주고(else), 이미 선이 존재하는 경우 더 낮은 값을 저장해준다.

 

<계산과 출력>

 

distances = dijkstra(graph, K)

for i in distances:
    print(distances[i] if distances[i] != float('INF') else 'INF')

 

함수의 최단 거리 반환 값을 distances에 대입한다.

경로가 존재 하지 않는 경우 inf 소문자로 출력되므로 대문자 INF로 출력되게 한다.

 

<전체 코드>

 

import sys
import heapq


def dijkstra(graph, start):
    distances = {node: float('INF') for node in graph}
    distances[start] = 0
    queue = []
    heapq.heappush(queue, [distances[start], start])

    while queue:
        current_distance, node = heapq.heappop(queue)

        if distances[node] < current_distance:
            continue

        for new_node, new_distance in graph[node].items():
            distance = new_distance+current_distance
            if distance < distances[new_node]:
                distances[new_node] = distance
                heapq.heappush(queue, [distance, new_node])

    return distances


V, E = map(int, sys.stdin.readline().split())
K = int(sys.stdin.readline())

graph = {}
for i in range(1, V+1):
    graph[i] = {}

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

    if v in graph[u]:
        if graph[u][v] > w:
            graph[u][v] = w
    else:
        graph[u][v] = w

distances = dijkstra(graph, K)

for i in distances:
    print(distances[i] if distances[i] != float('INF') else 'INF')