Baekjoon 1865. 웜홀 / Python

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

알고리즘 설명

문제 설명

주어진 그래프에 음수 가중치가 있는지를 확인하여 시간 여행이 가능한지(즉, 음수 사이클이 존재하는지) 판단하는 문제

코드 설명

모든 노드에서 출발하여 음수 사이클이 존재하는지 확인한다.
dp 배열을 사용하여 각 노드까지의 최단 거리를 저장하고, 그래프의 모든 간선을 N - 1번 반복하여 최단 거리를 갱신한다.
N번째 반복에서도 최단 거리가 갱신되면 음수 사이클이 존재한다고 판단한다.

import sys
input = sys.stdin.readline
inf = 1e10  

# 벨만-포드 알고리즘을 사용하여 음수 사이클을 탐지하는 함수
def sol(start=1):
    # 각 노드까지의 최단 거리를 저장하는 dp 배열
    dp = [inf for _ in range(n + 1)]
    dp[start] = 0  # 시작 노드의 거리는 0으로 설정
    
    # n번 반복하여 최단 거리를 갱신
    for idx in range(n):
        # 각 시작 지점에 대해
        for s in start_li:
            # 각 인접 노드와 그 간선의 가중치에 대해
            for e, t in arr[s]:
                # 최단 거리를 갱신하는 경우
                if dp[e] > dp[s] + t:
                    # n번째 반복에서도 갱신이 일어나면 음수 사이클이 존재
                    if idx == n - 1:
                        return True
                    dp[e] = dp[s] + t
    return False  # 음수 사이클이 없으면 False 반환


T = int(input())
for _ in range(T):
    # n, m, w (지점 수, 도로 수, 웜홀 수)
    n, m, w = map(int, input().split())
    
    arr = [[] for _ in range(n + 1)]
    start_li = set()  

    # 도로 정보
    for _ in range(m):
        s, e, t = map(int, input().split())
        arr[s].append([e, t])  # 무방향 도로
        arr[e].append([s, t])
        start_li.add(s)
        start_li.add(e)
    
    # 웜홀 정보 입력
    for _ in range(w):
        s, e, t = map(int, input().split())
        arr[s].append([e, -t])  # 단방향 웜홀, 시간이 줄어드는 경우
        start_li.add(s)
    
    start_li = list(start_li)

    res = sol()
    
    print(['NO', 'YES'][res])