1865. 웜홀 (Python)

1865번: 웜홀

 

 

풀이

 벨만 포드 알고리즘을 사용하여 음의 사이클이 존재하는지를 확인하면 된다. 벨만 포드 알고리즘은 링크참고하자.

 

 문제에는 시작 정점이 따로 주어져 있지 않다. 그렇기 때문에 웜홀의 정보를 입력받을 때 도착 지점에 대한 리스트를 Set을 활용하여 저장하였다. 웜홀의 정보를 입력받을 때 S에서 E로 웜홀을 통해서 가면 시간이 줄어들기 때문에 도착 지점인 E가 음의 사이클을 가질 수 있는 후보 지점이라고 생각했다. 

 이렇게 후보 지점들을 시작 지점으로 설정하여 벨만 포드 알고리즘을 수행하여 음의 사이클이 존재하는지를 체크하여 이에 따라 YES 또는 NO를 출력하도록 하였다.

 

 문제는 이렇게 구현하면 벨만 포드 알고리즘을 여러 번 수행할 수 있기 때문에 시간이 오래 걸린다는 단점이 있다.

 

 단 한 번의 벨만 포드 알고리즘 수행을 통하여 답을 찾을 수도 있다. 이에 대해서는 문제의 질문 게시판에 잘 설명되어 있는 게시글이 있으니 링크를 참고하자. 링크를 참고하여 임의의 시작 정점을 사용하여 벨만 포드 알고리즘을 한 번만 수행하면 시간을 줄일 수 있다.

 

코드

개선전 코드

import sys

input = sys.stdin.readline
INF = float('inf')


# 벨만 포드 알고리즘
# True: 음의 사이클 존재, False: 음의 사이클이 존재하지 않음
def bf(start):
    dist[start] = 0

    for i in range(N):
        for cur in range(1, N+1):
            for nextNode, cost in graph[cur]:
                if dist[cur] != INF and dist[nextNode] > dist[cur] + cost:
                    dist[nextNode] = dist[cur] + cost

                    if i == N-1:  # N번째 수행에서 최단 거리가 갱신 됐다면 음의 사이클이 존재
                        return True

    return False


if __name__ == '__main__':
    TC = int(input())  # 테스트 케이스 개수

    for _ in range(TC):
        N, M, W = map(int, input().split())  # N개의 지점, M개의 도로, W개의 웜홀
        graph = [[] for _ in range(N+1)]  # 도로와 웜홀의 정보
        V = set()  # 음의 사이클을 가질 수 있는 후보 지점들

        # 도로 정보 입력
        for _ in range(M):
            S, E, T = map(int, input().split())
            graph[S].append((E, T))
            graph[E].append((S, T))

        # 웜홀 정보 입력
        for _ in range(W):
            S, E, T = map(int, input().split())
            graph[S].append((E, -T))  # 웜홀의 경우 시간이 줄어드므로 시간을 음수 값으로 설정
            V.add(E)

        flag = False
        for start in list(V):  # 음의 사이클 후보 지점들에 대하여 벨만 포드 알고리즘 수행
            dist = [INF for _ in range(N+1)]
            flag = bf(start)

            if flag:  # 음의 사이클을 가지는 경우를 찾으면 break
                break

        if flag:
            print('YES')
        else:
            print('NO')

 

개선된 코드

import sys

input = sys.stdin.readline
INF = 10 ** 11  # ** float('inf')에서 10 ** 11로 변경 ** 


# 벨만 포드 알고리즘
# True: 음의 사이클 존재, False: 음의 사이클이 존재하지 않음
def bf(start):
    dist[start] = 0

    for i in range(N):
        for cur in range(1, N+1):
            for nextNode, cost in graph[cur]:
                # ** dist[cur] != INF 조건 삭제 **
                if dist[nextNode] > dist[cur] + cost:
                    dist[nextNode] = dist[cur] + cost

                    if i == N-1:  # N번째 수행에서 최단 거리가 갱신 됐다면 음의 사이클이 존재
                        return True

    return False


if __name__ == '__main__':
    TC = int(input())  # 테스트 케이스 개수

    for _ in range(TC):
        N, M, W = map(int, input().split())  # N개의 지점, M개의 도로, W개의 웜홀
        dist = [INF for _ in range(N+1)]
        graph = [[] for _ in range(N+1)]  # 도로와 웜홀의 정보

        # 도로 정보 입력
        for _ in range(M):
            S, E, T = map(int, input().split())
            graph[S].append((E, T))
            graph[E].append((S, T))

        # 웜홀 정보 입력
        for _ in range(W):
            S, E, T = map(int, input().split())
            graph[S].append((E, -T))  # 웜홀의 경우 시간이 줄어드므로 시간을 음수 값으로 설정

        if bf(1):
            print('YES')
        else:
            print('NO')