[알고리즘] 백준 1956 '운동' 풀이 _ 파이썬

백준 1956 운동
https://www.acmicpc.net/problem/1956
난이도 : 골드 4
유형 : 그래프 (플로이드 와샬/다익스트라)

문제

=> 방향성 있는 그래프에서 최소 비용으로 사이클을 돌 수 있는지 구하는 문제

풀이

1. 플로이드 와샬

경유지를 기준으로 경유지를 들렀을 경우와 아닌 경우를 비교하면서 최소 비용 루트를 찾는 플로이드 와샬 알고리즘을 사용해서 풀어봤다.

V, E = map(int, input().split())
#거리를 저장할 matrix
matrix = [[int(1e9)] * (V+1) for _ in range(V+1)]
for _ in range(E):
    x, y, c = map(int, input().split())
    matrix[x][y] = c

#경유지 k, 출발지 i, 도착지 j 로 3중 for문을 돌면서
for k in range(1, V+1):
    for i in range(1, V+1):
        for j in range(1, V+1):
            # i->j 가 빠른지 i->k->j가 빠른지를 검사한다.
            if matrix[i][k] + matrix[k][j] < matrix[i][j]:
                matrix[i][j] = matrix[i][k] + matrix[k][j]

answer = 1e9
for i in range(1, V+1):
   #사이클은 결국 출발지와 도착지가 같은 경우이므로 i->i를 확인
    answer = min(answer, matrix[i][i])

#최소값이 없으면 -1, 있으면 최소값을 출력
if answer == 1e9:
    print(-1)
else:
    print(answer)

플로이드 와샬을 사용하면 아무래도 3중 for문을 돌다보니 시간이 많이 소요된다.
이 문제에서도 python3으로 제출하면 시간초과, pypy로만 정답 인정이 됐다.

시간을 줄일 방법이 없을까 하고 백준 정답 코드를 구경하다가 변형된 다익스트라 알고리즘을 사용한 코드를 보고 참고해서 풀어봤다.

2. 다익스트라 (변형)

다익스트라 알고리즘은 heapq를 사용해서 가장 짧은 거리의 경로부터 검사하면서 이 경로를 거쳐서 갈 수 있는 노드까지의 거리를 계산하는 방식이다.

일반적으로 다익스트라는 하나의 노드에서 다른 노드까지 갈 수 있는 방법을 계산하는 알고리즘으로 시작점을 정한 뒤 1차원 배열에 다른 노드까지의 거리값을 저장하면서 진행되는데, 이를 약간 변형시켜 for문을 돌면서 시작점을 정해주고, 2차원 배열에 [시작점][도착점] 형태로 거리값을 저장해주는 방식으로 구현했다.

import heapq as hq

V, E = map(int, input().split())
graph = [[] for _ in range(V+1)]
#거리를 저장할 배열을 2차원으로
dist = [[1e9] * (V+1) for _ in range(V+1)]

heap = []
for _ in range(E):
    x, y, c = map(int, input().split())
    graph[x].append([c, y])
    dist[x][y] = c
    #heap에도 비용, 출발지, 도착지 3개의 값을 넣어준다.
    #유효한 경로 값을 다 넣어줌
    hq.heappush(heap, [c, x, y])


while heap:
    #최소 비용의 경로를 먼저 뽑아주고 (d:비용, s:출발, g:도착)
    d, s, g = hq.heappop(heap)
    #출발지와 도착지가 같다면 사이클!
    #heap을 이용하기 때문에 처음 나온 사이클이 가장 비용이 작은 사이클이므로 break 해버려도 됨! -> 여기서 시간이 굉장히 절약되는 듯 
    if s == g:
        print(d)
        break
    #d 값이 이미 저장된 비용보다 크다면 넘겨버림
    if dist[s][g] < d:
        continue
        
    #g에서 갈 수 있는 노드들을 검사
    for nd, ng in graph[g]:
    	# s->g->ng로 가는 비용
        new_d = d + nd
        # s->g->ng로 가는게 s->ng보다 빠르다면 값 갱신해주고 heap에 넣어줌
        if new_d < dist[s][ng]:
            dist[s][ng] = new_d
            hq.heappush(heap, [new_d, s, ng])
else:
    #heap 다 돌았는데 없다면 -1
    print(-1)

이 방식으로 풀었더니 python3으로도 1000ms 대로 수월하게 통과! 야호! 🎉
아무래도 heapq를 이용하면서 짧은 거리부터 검사해주고 그 과정에서 사이클을 찾으면 반복을 끝내버리기 때문에 시간을 대폭 줄일 수 있는 것 같다.

다익스트라 알고리즘이 아직 백지에서 구현할 만큼 익숙하지 않은 것 같다. 더 연습하기!😙