(Python) 백준 1865

백준 1865

시간 제한	메모리 제한	제출	정답	맞힌 사람	정답 비율
2 초	128 MB	47474	11143	6936	21.574%

문제

때는 2020년, 백준이는 월드나라의 한 국민이다. 월드나라에는 N개의 지점이 있고 N개의 지점 사이에는 M개의 도로와 W개의 웜홀이 있다. (단 도로는 방향이 없으며 웜홀은 방향이 있다.) 웜홀은 시작 위치에서 도착 위치로 가는 하나의 경로인데, 특이하게도 도착을 하게 되면 시작을 하였을 때보다 시간이 뒤로 가게 된다. 웜홀 내에서는 시계가 거꾸로 간다고 생각하여도 좋다.

시간 여행을 매우 좋아하는 백준이는 한 가지 궁금증에 빠졌다. 한 지점에서 출발을 하여서 시간여행을 하기 시작하여 다시 출발을 하였던 위치로 돌아왔을 때, 출발을 하였을 때보다 시간이 되돌아가 있는 경우가 있는지 없는지 궁금해졌다. 여러분은 백준이를 도와 이런 일이 가능한지 불가능한지 구하는 프로그램을 작성하여라.

입력

첫 번째 줄에는 테스트케이스의 개수 TC(1 ≤ TC ≤ 5)가 주어진다. 그리고 두 번째 줄부터 TC개의 테스트케이스가 차례로 주어지는데 각 테스트케이스의 첫 번째 줄에는 지점의 수 N(1 ≤ N ≤ 500), 도로의 개수 M(1 ≤ M ≤ 2500), 웜홀의 개수 W(1 ≤ W ≤ 200)이 주어진다. 그리고 두 번째 줄부터 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까지 자연수로 중복 없이 매겨져 있다.

출력

TC개의 줄에 걸쳐서 만약에 시간이 줄어들면서 출발 위치로 돌아오는 것이 가능하면 YES, 불가능하면 NO를 출력한다.

---

답안

import sys
import math
from collections import defaultdict

def initialize_variables():
    num_vertices, num_roads, num_wormholes = map(
        int, sys.stdin.readline().split())
    edges = defaultdict(lambda: math.inf)
    is_negative_cycle_detected = False
    # we are interested in whether the graph includes any negative cycles or not,
    # so we include edges with minimum weight values only.
    # Edges with non-minimum weight values are all ignored.
    # road: non-directed edges. start->end and end->start trip is all valid.
    for _ in range(num_roads):
        start, end, time = map(int, sys.stdin.readline().split())
        edges[(start, end)] = min(edges[(start, end)], time)
        edges[(end, start)] = min(edges[(end, start)], time)
    # wormhole: directed edges. only start->end trip is valid.
    for _ in range(num_wormholes):
        start, end, time = map(int, sys.stdin.readline().split())
        if start == end:
            is_negative_cycle_detected = True
        edges[(start, end)] = min(edges[(start, end)], -time)
    return edges, num_vertices, is_negative_cycle_detected

def bellman_ford(edges, start, num_vertices):
    distances = [math.inf for _ in range(num_vertices+1)]
    distances[start] = 0
    for _ in range(num_vertices - 1):
        # if distances is not changed during one round, it means you can early-stop the algorithm
        is_distances_changed = False
        for (start, end), time in edges.items():
            if distances[start]+time < distances[end]:
                distances[end] = distances[start]+time
                is_distances_changed = True
        if not is_distances_changed:
            return distances
    return distances

def is_distances_changed_during_bellman_ford_nth_round(edges, distances):
    for (start, end), time in edges.items():
        if distances[start]+time < distances[end]:
            return True
    return False

def solve():
    edges, num_vertices, is_negative_cycle_detected = initialize_variables()
    # if start node and end node is the same on one of wormholes, it means negative cycle(s) exist(s) in graph
    if is_negative_cycle_detected:
        print('YES')
        return
    # In case of given graph not being connected graph,
    # if not all nodes are visited, then start bellman-ford algorithm again with different starting node.
    visited = [False for _ in range(num_vertices+1)]
    for start in range(1, num_vertices+1):
        if not visited[start]:
            # In order to find out whether negative cycles are included in graph,
            # apply bellman-ford algorithm, n-1 rounds.
            distances = bellman_ford(edges, start, num_vertices)
            for node_number, distance in enumerate(distances):
                if not math.isinf(distance):
                    visited[node_number] = True
            # apply bellman-ford algorithm, one round more.
            # if distance values stays still, it means there no exists negative cycles in graph. otherwise there exists.
            is_negative_cycle_detected = is_distances_changed_during_bellman_ford_nth_round(
                edges, distances)
            if is_negative_cycle_detected:
                break
    if is_negative_cycle_detected:
        print('YES')
    else:
        print('NO')

num_testcases = int(sys.stdin.readline())
for _ in range(num_testcases):
    solve()

이 답안은 ‘맞았습니다’를 받았으나, 성능 면에서 더 개선될 여지가 있습니다.

추후 데이터 추가 시 ‘틀렸습니다’를 받을 수도 있으니, 주의 바랍니다.

자세한 것은 아래 풀이의 링크를 참고하세요.

풀이

나는 이 문제를 풀기 위해, 벨만-포드 알고리즘을 사용하였다.

벨만 포드 알고리즘은 n-1번의 라운드로 이뤄지는데, 이때 벨만 포드 알고리즘을 n-1 라운드까지 돌린 뒤 추가적으로 한 라운드를 더 돌렸을 때 거리 값이 갱신되는지 확인해봄으로써 음의 사이클이 존재하는지 알 수 있기 때문이다.

맨 처음에는 벨만 포드 알고리즘의 시작점을 노드 1번으로만 고정하고 시작했었는데, 문제 조건에서 그래프가 연결 그래프(connected graph)라는 조건이 없었으므로 노드 1번이 아닌 다른 노드에서도 벨만 포드 알고리즘을 돌릴 필요가 있었다.

그런데 이렇게 한다면 최악의 경우 모든 노드에서 벨만 포드 알고리즘을 동작시켜야 하므로, 원래의 벨만 포드 알고리즘의 시간 복잡도인 O(VE)에 정점의 개수가 곱해져 O(V^2E)가 되어버리게 되었다.

이 경우 python3 기준으로 시간 초과 판정을 받게 되었으며, O(V^2E)보다 더 효율적인 다른 방법이 있을까 많은 고민을 하다가 그 방법이 생각나지 않아, 시간 복잡도는 그대로이지만 벨만 포드 알고리즘을 최적화함으로써 ‘맞았습니다’ 판정을 받을 수 있었다.

‘맞았습니다’ 판정을 받았음에도 아주 찝찝한 느낌이 들어 다른 사람들의 풀이와 해설을 살펴봤는데, 아래와 같이 좋은 글을 찾을 수 있었다.

https://www.acmicpc.net/board/view/72995

이 글을 보며 내 코드가 이런 문제가 있었구나 하면서 크게 배울 수 있었고, 벨만 포드 알고리즘에 대해 더 깊이 이해하게 된 것 같다. 여러분도 꼭 이 글을 읽어보시길 추천한다.

아래부터는 내 코드가 어떤 흐름을 갖는지 설명해보도록 하겠다.

각 테스트케이스에 대해, 위와 같은 solve() 함수를 실행한다.

solve() 함수에서는 입력을 받아 필요한 변수들을 초기화한 뒤, 벨만 포드 알고리즘을 n번의 라운드 수행한 후 거리 값이 갱신되었는지 살펴본 뒤 이를 통해 음의 사이클이 존재하는지 판단한다.

def initialize_variables():
    num_vertices, num_roads, num_wormholes = map(
        int, sys.stdin.readline().split())
    edges = defaultdict(lambda: math.inf)
    is_negative_cycle_detected = False
    # we are interested in whether the graph includes any negative cycles or not,
    # so we include edges with minimum weight values only.
    # Edges with non-minimum weight values are all ignored.
    # road: non-directed edges. start->end and end->start trip is all valid.
    for _ in range(num_roads):
        start, end, time = map(int, sys.stdin.readline().split())
        edges[(start, end)] = min(edges[(start, end)], time)
        edges[(end, start)] = min(edges[(end, start)], time)
    # wormhole: directed edges. only start->end trip is valid.
    for _ in range(num_wormholes):
        start, end, time = map(int, sys.stdin.readline().split())
        if start == end:
            is_negative_cycle_detected = True
        edges[(start, end)] = min(edges[(start, end)], -time)
    return edges, num_vertices, is_negative_cycle_detected

먼저 stdin 입력을 받아 변수를 초기화하는 부분이다.

이후 벨만 포드 알고리즘을 쉽게 구현하기 위해, 그래프의 경우에는 edge를 모아둔 딕셔너리로 표현했다. edge는 {(시작점, 끝점): 가중치 }로 이뤄져 있다.

문제에서 주어진 ‘도로’의 경우 시작점→끝점, 끝점→시작점으로 가중치 t의 시간을 사용하여 이동할 수 있으므로, 그래프에 두 경우를 모두 추가해줬다.

‘웜홀’의 경우 방향이 존재하므로 시작점→끝점만을 가중치 -t의 시간을 사용하여(거꾸로 거슬러) 이동할 수 있도록 그래프에 추가해줬다.

이때, 우리는 그래프에 음의 사이클이 존재하는지에만 관심을 가지므로, 일부러 높은 값의 가중치를 가지는 간선을 이용할 필요가 없다. 따라서 동일한 시작점과 끝점을 가지는 간선이 둘 이상일 경우 그중 가장 낮은 가중치를 가지는 간선만을 사용하기로 했다.

그리고, 만약 시작점과 끝점이 같은 정점인 경우도 고려해줬는데, ‘도로’의 경우에는 시작점과 끝점이 같은 경우를 무시할 수 있지만 ‘웜홀’의 경우에는 시작점과 끝점이 하나의 정점으로 같은 경우에는 항상 음의 사이클이 존재하게 된다.

문제에서 시작점과 끝점이 같은 정점이 아니라는 조건이 없었으므로, 이와 같이 예외 상황을 처리해주었다.

# In case of given graph not being connected graph,
# if not all nodes are visited, then start bellman-ford algorithm again with different starting node.
visited = [False for _ in range(num_vertices+1)]
for start in range(1, num_vertices+1):
    if not visited[start]:
        # In order to find out whether negative cycles are included in graph,
        # apply bellman-ford algorithm, n-1 rounds.
        distances = bellman_ford(edges, start, num_vertices)
        for node_number, distance in enumerate(distances):
            if not math.isinf(distance):
                visited[node_number] = True
        # apply bellman-ford algorithm, one round more.
        # if distance values stays still, it means there no exists negative cycles in graph. otherwise there exists.
        is_negative_cycle_detected = is_distances_changed_during_bellman_ford_nth_round(
            edges, distances)
        if is_negative_cycle_detected:
            break

위는 solve() 함수의 일부분을 가져온 것이다.

정점의 개수를 V라고 할 때, 벨만 포드 알고리즘의 시작점을 V번을 초과해 정하는 것을 막기 위해, visited라는 boolean 배열을 두었다. 만약 벨만 포드 알고리즘을 돌린 뒤에도 여전히 math.inf로 남아 있으면 그 정점에서 벨만 포드 알고리즘을 다시 돌린다.

(사실 위 링크의 글을 읽어 보면 알겠지만, 더 좋은 방법이 존재한다)

벨만 포드 알고리즘을 n번째 라운드까지 돌린 후 distances에 갱신이 이뤄졌다면 is_negative_cycle_detected를 True로 만든다.

이를 이용해 정답을 stdout으로 출력하면 된다.

이 코드가 개선될 수 있는 부분은,

• math.inf가 아닌 특정한 정수값을 사용한다면, visited 배열을 쓰지 않고도 노드에 방문했는지 확인 가능하다. 이때 우리는 벨만 포드 알고리즘의 결과값인 최단 경로가 어떤 값이 되는지에는 전혀 관심이 없고, 그 값이 바뀌는지에만 관심을 가지므로 어떠한 정수값을 사용해도 좋다.
• 하나의 임의의 정점을 만든 뒤 모든 정점을 그 정점에 연결하고, 그 가중치를 0으로 한다면, 연결 그래프가 아닐 경우를 대비해, 정점의 개수만큼 벨만 포드 알고리즘을 반복하여 돌릴 필요가 없다. 이렇게 한다면 O(VE)만에 이 문제를 해결할 수 있다.

위와 같다.

math.inf에서 어떠한 숫자값을 빼면 그대로 math.inf가 된다는 사실과, 그 사실 때문에 내가 짠 코드의 결과에 큰 영향을 줄 수 있다는 점은 놀라웠다.

다음부터는 아주 큰 값을 무조건 math.inf로 설정하지 않고, 그 side effect에 주의하여 알고리즘을 짜야겠다는 생각이 들었다.