[백준 11404번] 플로이드 (파이썬 풀이)

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

문제

문제 이해

Dijkstra로 풀려고 했는데, 시간초과가 났다. 문제의 이름처럼 플로이드 - 위셜 알고리즘을 써서 풀어야 하나보다.

chatGPT 에게 그게 뭐냐고 물어봤다. 대답은 잘한다.

코드로 나타내면 이해가 쉽다. 모든 노드 쌍간의 최단경로를 한번의 구할 수 있고 음의 가중치도 가능하다는 장점이 있다.

graph는 어떻게 나타내느냐 하면,

graph = [
    [0, 1, 1, inf, inf],
    [1, 0, 1, 1, inf],
    [1, 1, 0, inf, 1],
    [inf, 1, inf, 0, 1],
    [inf, inf, 1, 1, 0]
]

이런식으로 나타낸다고 한다. 문제에 적용해 보자.

 

문제 풀이

import sys
input = sys.stdin.readline

def floyd_warshall(graph):
    dist = graph.copy()

    # 모든 노드 쌍간의 최단경로를 구하는 알고리즘
    for k in range(n):
        for i in range(n):
            for j in range(n):
                dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])

    # 여전히 무한대로 남아있는 것을 0으로 바꿔준다. (문제에 그렇게 하라고 나와있다)     
    for i in range(n):
        for j in range(n):
            if dist[i][j] == float('inf'):
                dist[i][j] = 0
    return dist


n = int(input())
m = int(input())
graph = [[float('inf')]*(n) for _ in range(n)]
for i in range(n):
    graph[i][i] = 0

# graph 만드는 작업    
for _ in range(m):
    u, v, e = map(int, input().split())
    if e < graph[u-1][v-1]:
        graph[u-1][v-1] = e

# 출력 작업
for arr in floyd_warshall(graph):
    print(' '.join(map(str,arr)))

문제 풀이후 생각 정리

Dijkstra로 풀었던게 시간초과가 걸리니까 당황했다. 노드의 수가 적은 경우, 그리고 모든 최단 경로를 구하는 경우는 플로이드 - 워셜 알고리즘이 더 빠르다는 것을 알았다. 잊지 않고 잘 숙지 해야겠다.

아래는 시간초과 난 다익스트라 코드다,,

import sys
import heapq
input = sys.stdin.readline

def dijkstra(start):
    distance = [float('inf')] * (n+1)
    heap = [(0, start)]
    distance[start] = 0
    while heap:
        d, no = heapq.heappop(heap)
        for node, dis in edge[no]:
            D = dis + distance[no]
            if D < distance[node]:
                distance[node] = D
                heapq.heappush(heap, (D, node))
    return distance[1:]

n = int(input())
m = int(input())
edge = [ [] for _ in range(n+1)]

for _ in range(m):
    u, v, e = map(int, input().split())
    edge[u].append((v,e))
for i in range(1, n+1):
    print(' '.join(map(str,dijkstra(i))))