[백준 BOJ_1865] 웜홀 Python 풀이

출처: 백준 온라인 저지

문제

풀이

벨만-포드 (Bellman-Ford) 알고리즘을 이용하여 풀어주었습니다.

우선 벨만-포드 알고리즘은 가중치가 음수일 때에도 사용가능한 최단거리 알고리즘입니다. 벨만-포드 알고리즘에서는 모든 간선을 정점의 수만큼 반복하여 비교해주는데, 여기서 정점의 수만큼 반복했음에도 아직도 갱신되는 최단거리가 있다면, 그 부분은 최단거리를 무한히 줄이는 음수 간선의 순환, 음수 사이클이 존재한다는 뜻이 됩니다. 이 문제에서는 그 음수 사이클이 있는 지에 대한 여부를 구하는 문제가 됩니다.

그래서 dist를 가중치인 T가 같는 최대의 숫자 10,000에서 1을 더한 값으로 초기화해줍니다. math.inf나, float(“inf”)를 썼을 때에는 아무리 음수를 더해주어도 inf로 인식하기 때문에 갱신되는지 여부를 판단할 수 없게 되어버리기 때문에 이번 문제에서는 쓰지 않는게 좋습니다.

그 후 정점의 수인 N번만큼 모든 간선들을 확인해주어 그 정점 S를 거쳐 E로 가는 것이 정점 E까지 가는 최단거리보다 적다면 갱신해줍니다. 만약 N번째에도 갱신이 되고 있다면 음수 사이클이 존재한다는 뜻이기 때문에 바로 1을 return해주고, 그렇지 않고 무사히 for loop을 끝마쳤다면 0을 return해줍니다.

입력값을 받을 때, 도로는 방향이 없기 때문에 edges에 양방향으로 두 번 추가해주고, 웜홀은 방향이 있기 때문에 단방향으로 한 번만 추가해줍니다. 또한, 따로 웜홀의 가중치가 음수로 표시되어있지 않기 때문에 웜홀의 가중치는 음수로 추가해주어야합니다.

마지막에는 bellman-ford 함수의 결과에 따라 알맞은 출력값을 출력해줍니다.

코드

import sys
input = sys.stdin.readline
def bellman_ford(start, N, edges):
  dist = [10001 for _ in range(N+1)]
  dist[start] = 0
  for i in range(N):
    for S, E, T in edges:
      if dist[S] + T < dist[E]:
        dist[E] = dist[S] + T
        if i == N-1:
          return 1
  return 0
TC = int(input())
for _ in range(TC):
  N, M, W = map(int, input().split())
  edges = []
  
  for _ in range(M):
    S, E, T = map(int, input().split())
    edges.append((S, E, T))
    edges.append((E, S, T))
  for _ in range(W):
    S, E, T = map(int, input().split())
    edges.append((S, E, -T))
  
  print("YES" if bellman_ford(1, N, edges) else "NO")