프로그래머스 / 1865 / 웜홀 / python

처음시도 코드

import sys
input = sys.stdin.readline

T = int(input())

def bellman(start):
    dist = [1e9] * (n+1)
    dist[start] = 0

    for i in range(1, n+1):
        for j in range(len(graph)):
            # print(dist)
            node, next_node, dis = graph[j]

            if dist[node] != 1e9 and dist[next_node] > dist[node] + dis:
                dist[next_node] = dist[node] + dis

                if i == n:
                    return True
    
    return False

for _ in range(T):
    n, m ,w = map(int, input().split())

    # dist = [1e9] * (n+1)
    graph = []
    for _ in range(m):
        a, b, g = map(int, input().split())
        graph.append([a,b,g])
        graph.append([b,a,g])
    for _ in range(w):
        a, b, g = map(int, input().split())
        graph.append([a,b,-g])

    for i in range(1, n+1):
        if bellman(i):
            print("YES")
            break
    else:
        print("NO")

벨만 포드 알고리즘으로 시도
주어진 출발 지점이 없기 때문에 for 문을 통해 모든 노드를 시작점으로해서 탐색

시간초과가 발생한다....

시간 복잡도를 생각해보면 모든 노드에 대해서 탐색을 시작하므로 (n x 2m x w) x n이 된다. -> 500 x 500 x 5000 x 200 = 250,000,000,000....
엄청 크다....

이 문제의 경우 루프가 있는지 확인만 하면 되는 문제이다. 그렇기 때문에 시작점을 기준으로 거리를 구하는 로직을 빼고 dist를 통해 음수 루프가 있는지만 확인하면 된다.

그러기 위한 방법으로는 for 문을 돌때 dist[node] != 1e9 코드를 지워주면된다.

이 코드는 시작점으로 부터 거리를 구하기 위해서 해당 dist로 부터 이어지는 노드들을 구하기 위함이다. 따라서 위 코드를 없애주면 dist는 시작점으로 부터 각 노드의 거리 값을 저장하진 않고 루프가 있는지 확인하는 로직이 된다.

이렇게 하면 시간복잡도는 500 x 5000 x 200로 많이 줄어든다.

아래는 수정한 코드이다.

import sys
input = sys.stdin.readline

T = int(input())

def bellman(start):
    dist = [1e9] * (n+1)
    dist[start] = 0

    for i in range(1, n+1):
        for a,b,c in graph:
            # print(dist)
            node, next_node, dis = a,b,c

            if dist[next_node] > dist[node] + dis:
                dist[next_node] = dist[node] + dis

                if i == n:
                    return True
    
    return False

for _ in range(T):
    n, m ,w = map(int, input().split())

    # dist = [1e9] * (n+1)
    graph = []
    for _ in range(m):
        a, b, g = map(int, input().split())
        graph.append([a,b,g])
        graph.append([b,a,g])
    for _ in range(w):
        a, b, g = map(int, input().split())
        graph.append([a,b,-g])

    if bellman(1):
        print("YES")
    else:
        print("NO")

---

이 문제를 통해 벨만 포드 알고리즘을 다시한번 기억하고 또한 자세히 생각할 수 있게 되었다.