[골드 3] 백준 16437 - 양 구출 작전 (파이썬)

https://www.acmicpc.net/problem/16437

풀이

• 각 섬에는 양(S) 또는 늑대(W)가 있으며,
• 한 섬에 있는 동물들은 연결된 부모 섬(루트 방향)으로 이동한다.
• 만약 늑대 수가 양보다 많으면 그만큼 양이 잡아먹힌다.
• 목표: 1번 섬(루트)에 도착할 수 있는 양의 최대 수 구하기

입력처리는 다음과 같이 지정한다.

N = int(input())  # 섬의 개수
A = [0] * N       # 각 노드에 있는 양 또는 늑대 수 (양은 양수, 늑대는 음수)
child = [[] for _ in range(N)]  # 트리 구조 저장 (부모 -> 자식)

for i in range(1, N) :
    t, a, p = input().split()
    a = int(a)
    p = int(p) - 1
    if t == 'S' :
        A[i] = a        # 양이면 그대로 저장
    else :
        A[i] = -a       # 늑대이면 음수로 저장
    child[p].append(i)  # 부모 노드에 자식 추가

입력에서 2번째 노드부터 시작하고, 부모-자식 구조를 트리로 구성한다.

D = [0] * N  # 각 노드에서 구해올 수 있는 양의 수
 
def dfs(u) :
    D[u] = A[u]  # 현재 노드의 양 또는 늑대 수
 
    for v in child[u] :
        dfs(v)           # 자식 노드 먼저 탐색 (후위 순회)
        D[u] += D[v]     # 자식 노드로부터 구해온 양 수 더함
 
    if D[u] < 0 :
        D[u] = 0         # 늑대가 많아서 잡아먹으면 양 수는 0

DFS + DP를 수행하며, DFS를 하면서 자식 노드들로 부터 구해올 수 있는 최대 양을 누적한다.

이떄 누적 후 0보자 작으면 양이 다 잡아먹힌 것이므로 0으로 초기화한다.

마지막에 dfs(0)을 통해서 루트 노드에서 시작해 전체 DFS를 돌린 후, 루트에서 구할 수 있는 양의 총합을 출력한다.

코드

# 백준 16437 - 양 구출 작전
# 분류 : 다이나믹 프로그래밍
 
import sys
sys.setrecursionlimit(12345678)
 
N = int(input())
A = [0] * N
child = [[] for _ in range(N)]
 
for i in range(1, N) :
    t, a, p = input().split()
    a = int(a)
    p = int(p) - 1
    if t == 'S' :
        A[i] = a
    else :
        A[i] -= a
    
    child[p].append(i)
 
D = [0] * N
 
def dfs(u) :
    D[u] = A[u]
    for v in child[u] :
        dfs(v)
        D[u] += D[v]
    
    if D[u] < 0 :
        D[u] = 0
 
dfs(0)
print(D[0])