[백준 1865] 웜홀 - 플로이드 워셜 풀이

문제 출처 : https://www.acmicpc.net/problem/1865

1865번: 웜홀

이 문제는 벨만 포드로 풀어야 하는 문제 입니다. 벨만 포드 풀이법은 다빈치코딩 알고리즘 책에 작성 하였습니다.

https://wikidocs.net/225180

02. 웜홀[백준 1865]

 

정상적인 풀이법은 위키독스에 적어 놓았지만 이 문제를 풀기 위해 고생했던 것을 생각해서 플로이드 워셜 풀이법도 이곳에 작성해 놓습니다.

최단 거리 알고리즘 정하는 법

여러가지 최단 거리 알고리즘 중 어떤 알고리즘을 선택해야 하는지 정하는 방법부터 알아보겠습니다. 일단 기본적으로 최단 거리를 구하는 경우에는 BFS를 떠올릴 수 있습니다. BFS 알고리즘 역시 최단 거리를 구할 수 있지만 전제가 정점과 정점 사이의 거리가 모두 같다는 전제가 있습니다. 거리의 길이가 모두 다른 경우에는 BFS를 사용할 수 없습니다. 

 

거리가 모두 다른 경우에 사용하는 알고리즘은 다익스트라 알고리즘을 사용합니다.  BFS를 응용한 알고리즘으로 왠만한 최단 거리는 다익스트라로 거의 구할 수 있습니다.

 

다익스트라 알고리즘을 사용하면 거의 모든 최단 경로를 구할 수 있지만 못 구하는 경우가 있습니다. 바로 거리가 음수인 경우 입니다. 거리가 음수인 경우 없어 말이 안된다고 생각할 수 있지만 다양한 문제들에서는 얼마든지 음수인 경우를 만들 수 있습니다. 그래서 타임머신을 타서 시간을 줄일다거나, 이 문제처럼 웜홀로 거리를 줄인다거나, 도시를 방문할 때 교통비를 사용하거나, 상금을 주는 방식으로 음수를 사용할 수 있는 문제들을 만들어 놓았습니다. 이렇게 음수가 있는 문제는 다익스트라가 아닌 벨만 포드 알고리즘을 사용합니다. 

 

벨만 포드 알고리즘이 음수가 있는 경우에 최단 거리를 찾지만 이 알고리즘은 시작 위치가 정해져 있습니다. 시작 위치가 정해지지 않았을 경우에는 플로이드 워셜 알고리즘을 사용합니다. 플로이드 워셜 알고리즘은 음수의 거리가 있어도 풀 수 있습니다.

 

단 벨만 포드 알고리즘이나 플로이드 워셜 알고리즘 모두 음의 사이클이 있는 경우 거리를 구할 수 없습니다. 음의 사이클이란 일정 구역을 돌때마다 시간이 점점 줄어드는 경우를 말합니다.

왼쪽 사이클 1 -> 2 -> 3의 경우 음수가 있기는 하지만 한 바퀴를 돌고 나면 거리가 3 + 1 + (-2) 로 2가 됩니다. 2바퀴를 돌고 나면 거리가 4가 됩니다. 반면에 오른쪽 사이클인 4 -> 5 -> 6 도 음수가 있습니다. 한 바퀴를 돌고나면 3 + 1 + (-5)로 -1이 됩니다. 2바퀴를 돌고나면 거리가 -2가 되고 여러바퀴를 돌면 돌수록 거리가 줄어들게 됩니다. 이런 식으로 음수로 계속 줄어드는 사이클을 음의 사이클이라고 합니다.

 

문제 풀이하기

이 문제에는 시작 위치가 정해지지 않았고, 음수 거리가 있기 때문에 플로이드 워셜이 제격이라 생각했습니다. 또한 정점이 500개이기 때문에 플로이드 워셜의 시간 복잡도가 $ O(V^3) $ 이기 때문에 반복횟수가 1억이 조금 넘고 시간이 2초이기 때문에 충분하다고 생각했습니다. 테스트 케이스 5번을 고려하지 않은 것이 문제였습니다. 어찌되었건 플로이드 워셜로 풀긴 풀었는데 계속 시간초과가 발생해서 벨만 포드로 변경을 했습니다. 하지만 여러 질문들을 보니 플로이드 워셜로 푼 사람들이 있어 오기가 들어 플로이드 워셜로 결국 풀게 되었습니다. 잘못된 풀이를 책에 쓰기는 애매해서 블로그에 공개합니다.

 

입력 받기

import sys
input = sys.stdin.readline

mii = lambda : map(int, input().split())
TC = int(input())

for _ in range(TC):
    N, M, W = mii()
    INF = N * 10000
    arr = [[INF] * (N + 1) for _ in range(N + 1)]
    for _ in range(M):
        s, e, t = mii()
        arr[s][e] = min(arr[s][e], t)
        arr[e][s] = min(arr[e][s], t)
    
    for _ in range(W):
        s, e, t = mii()
        arr[s][e] = min(arr[s][e], -t)

    floyd()

맨 처음 테스트 케이스 TC를 입력 받습니다. 다음으로 정점의 수 N, 거리 M, 웜홀 W를 입력 받습니다. 다음으로 플로이드 워셜 알고리즘은 인접 행렬 형태이기 때문에 2차 배열로 만들어 줍니다. 인접 행렬 형태이기 때문에 메모리 낭비가 심합니다. 인접 행렬에 대해 잘 모른다면 여기를 확인해 보시기 바랍니다.

 

다음으로 거리를 입력 받습니다. 도로는 양방향 도로이기 때문에 양쪽으로 입력을 받습니다. 이 때 문제에서 두 지점을 연결하는 도로가 한 개보다 많을 수 있다고 했기 때문에 거리의 최소값을 저장하도록 합니다.

 

웜홀은 단방향이기 때문에 입력 그대로 저장을 합니다. 웜홀은 시간이 음수이지만 입력받을 때에는 양수로 받기 때문에 마이너스를 붙여 음수로 바꿔줍니다. 웜홀 역시 여러개가 입력될 수 있기 때문에 최소값을 저장하도록 합니다.

 

마지막으로 floyd() 라는 함수를 실행하여 결과를 출력하도록 합니다. 그럼 floyd 함수에 대해 알아보겠습니다.

플로이드 워셜 알고리즘 함수

def floyd():
    for k in range(1, N + 1):
        for i in range(1, N + 1):
            if k == i:
                continue
            for j in range(1, N + 1):
                if k == j:
                    continue
                arr[i][j] = min(arr[i][j], arr[i][k] + arr[k][j])
    
    for i in range(1, N + 1):
        if arr[i][i] < 0:
            print("YES")
            return

    print("NO")
    return

플로이드 워셜 알고리즘은 정점들을 모두 돌면서 더 짧은 거리가 있으면 계속 갱신하는 방식으로 진행합니다. 이 때 음의 사이클이 있는지 확인하는 방법은 자기 자신으로 가는 값이 0보다 작은 경우 입니다. 위에서 arr[i][i] < 0 을 비교하는 이유가 음의 사이클이 있는지 확인하는 부분 입니다. 모든 정점들에 대해서 음의 사이클이 있는지 확인하여 있는 경우 YES를 출력하고 없는 경우 NO를 출력하도록 한 것입니다.

정리

이 문제는 플로이드 워셜이 아닌 벨만 포드로 푸는 것이 맞습니다. 시간 복잡도를 따져봐도 정점 500개에 테스트 케이스 5개면 반복 횟수가 6억이 넘어가기 때문에 2초안에 해결이 안되는 것이 정상입니다. 다만 이 문제의 경우 잘만 풀면 시간초과가 나지 않고 통과가 되기 때문에 풀이한 것이 아까워서 올려놓은 것입니다.

이것이 시간 초과를 뚫어볼려고 이리저리 시도한 영광의 상처 입니다. 정상적인 플로이드 워셜 풀이는 안되게 막아놓은것 같습니다.

 

결국 맞은 내역을 보면 벨만 포드 알고리즘을 사용했을 때는 400ms 이하로 해결 되었고, 4천ms, 6천ms가 걸린 것이 플로이드워셜로 푼것 입니다.

전체 코드 보기

마지막으로 전체코드를 확인해 보겠습니다. 전체 코드를 따로 꼭 넣는 이유는 파이썬의 특성상 들여쓰기가 중요한데 소스를 잘라서 올리면 어디서 들여쓰기를 해야하는지 애매해 하시는 분들이 많아서 입니다. 또한 복사해서 바로 쓰기에도 좋기 때문에 전체 코드를 꼭 추가하는 편입니다.

import sys
input = sys.stdin.readline

mii = lambda : map(int, input().split())

def floyd():
    for k in range(1, N + 1):
        for i in range(1, N + 1):
            if k == i:
                continue
            for j in range(1, N + 1):
                if k == j:
                    continue
                arr[i][j] = min(arr[i][j], arr[i][k] + arr[k][j])
    
    for i in range(1, N + 1):
        if arr[i][i] < 0:
            print("YES")
            return

    print("NO")
    return

TC = int(input())

for _ in range(TC):
    N, M, W = mii()
    INF = N * 10000
    arr = [[INF] * (N + 1) for _ in range(N + 1)]
    for _ in range(M):
        s, e, t = mii()
        arr[s][e] = min(arr[s][e], t)
        arr[e][s] = min(arr[e][s], t)
    
    for _ in range(W):
        s, e, t = mii()
        arr[s][e] = min(arr[s][e], -t)

    floyd()