[BOJ 16437] 양 구출 작전 (Python)

복습 요망!!

풀이

처음 union-find 알고리즘의 find 알고리즘을 이용하여 탐색하는 섬이 1번 섬으로 이동할 수 있는지 먼저 판단한 후, 탐색하는 섬에서부터 1번 섬으로 가는 경로가 존재한다면 1번 섬으로 가는 중에 양 - 늑대를 계산하면서 모든 경로를 전부 탐색하는 것으로 접근하였다. 하지만, 구현하는 것이 벅차 다른 사람의 풀이를 참고하여 풀었다. (이상하게 풀이만 보면 이해가 잘된다...)

아래 그림을 통해 문제를 정확하게 이해해보자. 심지어 문제에서 그림을 제시해줬다.

추가로 늑대는 섬을 이동할 수 없다고 나와있으므로 4번 섬의 늑대는 1번 섬으로 이동할 수 없다. 문제는 확실히 이해했으나 구현하기 어렵다면 코드를 복붙해서 직접 디버깅 돌려보자. (본인도 그렇게 했다...)

코드

import sys

# 시간초과 발생으로 인해 recursion 조정
sys.setrecursionlimit(10 ** 6)
input = sys.stdin.readline
N = int(input())
graph = [[0, []] for _ in range(N + 1)]
for i in range(2, N + 1):
    animal, animal_cnt, island = input().split()
    if animal == 'W':
        animal_cnt = -int(animal_cnt)
    graph[i][0] = int(animal_cnt)
    graph[int(island)][1].append(i)

# print(graph)


def dfs(depth):
    # 늑대 or 양 마리 수 저장 및 갱신
    answer = graph[depth][0]
    for g in graph[depth][1]:
        answer += dfs(g)

    # 섬을 이동할 때 섬에 늑대만 있거나 양보다 늑대가 많아서 양을 다 잡아 먹을 때
    if answer < 0:
        return 0
    else:
        return answer


print(dfs(1))