[BOJ] 백준 16437 양 구출 작전 python

import sys

sys.setrecursionlimit(1000000)
def solution():
    input = sys.stdin.readline
    N = int(input())
    #wolves[i] : i번째 섬에 있는 늑대 수
    wolves =[0 for _ in range(N+1)]
    
    #sheeps[i] : i번쨰 섬에 있는 양의 수=
    sheeps = {i:0 for i in range(1, N+1)}
    #트리
    tree =[[] for _ in range(N+1)]
    for i in range(2, N+1):
        t, a, p = input().split()
        a = int(a)
        p = int(p)
        if t == "W":
            wolves[i] = a
        else:
            sheeps[i] = a

        tree[p].append(i)
    
    def dfs(here):
        num_sheeps = sheeps[here]

        for there in tree[here]:
            num_sheeps += dfs(there) #서브트리의 양들을 모두 더해준다.

        if wolves[here] != 0: #만약 현재 섬에 늑대가 있고
            if num_sheeps < wolves[here]: #서브트리의 양보다 현재 섬의 늑대가 많다면
                wolves[here] -= num_sheeps #배부른 늑대들이 줄어든다.
                num_sheeps =0 #양들은 모두 먹힘
            else: #양이 더 많다면
                num_sheeps -= wolves[here] #양이 줄어들고
                wolves[here] =0 #배부른 늑대는 0 이제부턴 이 섬은 동료들의 희생으로 그냥 지나갈 수 있다.
                
        return num_sheeps
    
    print(dfs(1))
    

solution()


"""
9 
S 10 1
S 10 1
W 10 1
S 10 2
W 10 2
S 10 6
S 10 7
S 10 7

50
"""