[백준] 16437번 양 구출 작전 - 파이썬/트리

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

import sys
sys.setrecursionlimit(10**6)

input = sys.stdin.readline
N = int(input())

graph = [ [] for _ in range(N+1) ]
visited = [False] * (N+1)
animal = [0] * (N+1)
mari = [0] * (N+1)
dp = [0] * (N+1)

for i in range(N-1):

    t,a,p = input().split()

    a = int(a)
    p = int(p)

    animal[i+2]= t
    mari[i+2] = a

    graph[i+2].append(p)
    graph[p].append(i+2)


def dfs(x):
    visited[x]=True

    for i in graph[x]:
        if not visited[i]:
            dfs(i)
            dp[x]+=dp[i]

    if animal[x]=='W':
        dp[x] -= mari[x]
        if dp[x] < 0:
            dp[x]=0
    elif animal[x]=='S':
        dp[x] += mari[x]

    if len(graph[x])==1 and animal[x]=='S':
        dp[x]=mari[x] # 리프노드에서 dp값 설정

dfs(1)
print(dp[1])

트리+DP 유형의 문제이다.

비슷한 유형의 문제를 많이 풀어보다보니 어느정도 감이 잡혀서 쉽게 풀 수 있었다.

x 노드에서 dfs를 할때 자식노드를 방문하고 방문 후에는 자식 노드의 dp값을 더해준다.

그러고 나서 x 노드의 동물이 양이라면 dp[x] 에 더해주고, 늑대라면 dp[x]에서 빼준다.

그런데 양이 음수가 될 수 없으므로 dp[x]가 0보다 작으면 0으로 입력해준다.