[백준 1865번] 파이썬 - 웜홀

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

1865번: 웜홀

 

# 조건

• 월드나라에는 N개의 지점이 존재하고
  • N개의 지점 사이에는 M개의 도로와 W개의 웜홀
  • 도로는 무방향, 웜홀은 방향이 존재
• 웜홀은 시작 위치에서 도착 위치로 가는 하나의 경로인데, 도착을 하게 되면 시작 하였을 때보다 시간이 뒤로 가게 된다.
• 한 지점에서 출발하여 시간여행을 하기 시작하여서 다시 출발을 하였던 취로 돌아왔을 때,
• 출발을 하였을 때보다 시간이 되돌아가 있는 경우가 있는지 없는지 프로그램을 작성하여라.

입력

• 첫 번째 줄에는 테스트케이스의 개수 TC(1 ≤ TC ≤ 5)가 주어진다.
• 두 번째 줄부터 TC개의 테스트케이스가 차례로 주어지는데 각 테스트케이스의 첫 번째 줄에는 지점의 수 N(1 ≤ N ≤ 500), 도로의 개수 M(1 ≤ M ≤ 2500), 웜홀의 개수 W(1 ≤ W ≤ 200)이 주어진다.
• 각 TC의 두 번째 줄부터 M+1번째 줄에 도로의 정보가 주어지는데 각 도로의 정보는 S, E, T 세 정수로 주어진다.
  • S와 E는 연결된 지점의 번호, T는 이 도로를 통해 이동하는데 걸리는 시간을 의미한다.
  • 그리고 M+2번째 줄부터 M+W+1번째 줄까지 웜홀의 정보가 S, E, T 세 정수로 주어지는데 S는 시작 지점, E는 도착 지점, T는 줄어드는 시간을 의미한다. T는 10,000보다 작거나 같은 자연수 또는 0이다.
• 두 지점을 연결하는 도로가 한 개보다 많을 수도 있다. 지점의 번호는 1부터 N까지 자연수로 중복 없이 매겨져 있다.

출력

• 돌아오는 것이 가능하면 YES
• 아니라면 NO

# 접근 방법 및 Solution

• 문제 이해부터, 입력을 받아주는 것이 너무 헷갈리는 문제인 것 같다.
• 웜홀을 통과했을 경우 시간이 -가 되므로 다익스트라가 아닌 '벨만-포드' 알고리즘을 이용해주면 된다.
• 일반적인 벨만-포드의 경우 '음의 사이클'이 존재하면 안되므로 유무를 확인해준다.
• 이번 문제의 경우 모든 정점에서의 시작 -> 노드의 최소 거리를 구하며 돌아왔을 경우 -가 되어있으면 YES를 출력하는 문제인데
  • 위 방법처럼 모든 정점 확인시 시간초과 발생
• 따라서 기존의 dis[cur] != INF 조건
  • BFS 함수에서 주어진 인자를 시작지점으로 이어진 노드의 최소 거리를 구하기 위한 조건
• 을 제거 해주고, 아무 점에서 시작하여 n번째 루프를 돌 때 최단거리 테이블의 변화 유무를 통해 **음의 사이클**이 존재한다면, YES를 출력해주면 된다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  
INF = int(1e9)  
  
def bellman_ford(s):  
    global N  
    visited = [INF for _ in range(N+1)]  
    visited[s] = 0  
  
    # 음의 사이클 판별 위해 n번 반복  
    for i in range(N):  
        # 반복마다 모든 간선 확인  
        for start, end, dist in info:  
            # 지금 계산한 거리가, 기존 계산된 거리보다 작다면 테이블 갱신  
            if visited[start] + dist < visited[end]:  
                visited[end] = visited[start] + dist  
                # i가 n-1이고 테이블이 갱신되었다면 '음의 싸이클' 존재  
                if i == N-1:  
                    return True  
    return False  
T = int(input())  
for tc in range(T):  
    N, M, W = map(int, input().split())                 # 지점의 수, 도로의 수, 웜홀의 수  
    info = []                                           # 도로와 웜홀 정보  
    for _ in range(M):  
        s, e, t = map(int, input().split())  
        info.append((s,e,t))  
        info.append((e,s,t))                           # 도로의 경우 무방향이므로 양쪽 다 기록  
  
    for _ in range(W):  
        s, e, t = map(int, input().split())  
        info.append((s, e, -t))                          # 웜홀의 경우 음수 가중치이므로 - 붙여서 받아주기  
  
    flag = bellman_ford(1)  
    if flag:  
        print('YES')  
    else:  
        print('NO')

 

• 위 코드의 경우 출발점이 1로 지정되어있지만 통과가 된다.
• 하지만 문제에서 1로 고정되어있지 않기 때문에 아래와 같이 모든 지점에서 체크를 해주는 것이 올바른 코드인 것 같다.

 

import sys
input = sys.stdin.readline

def BellmanFord(n, edge, source):
    dist = [float('inf')]*(n+1); dist[source] = 0
    for i in range(n-1):
        for u,v,c in edge: dist[v] = min(dist[v], dist[u]+c)
    for u,v,c in edge:
        if dist[u]+c < dist[v]: return False
    return dist

for T in range(int(input())):
    n, m, w = map(int,input().split())
    edge = []
    for i in range(n):
        edge.append((0,i+1,0))
    for i in range(m):
        s, e, t = map(int,input().split())
        edge.append((s,e,t))
        edge.append((e,s,t))
    for i in range(w):
        s, e, t = map(int,input().split())
        edge.append((s,e,-t))
    print("YES" if not BellmanFord(n, edge, 0) else "NO")