문제 분석
단순한 DFS 문제로 판단했습니다. 다른 방식은 아직 생각나지 않았습니다.
어차피 문제의 자료구조가 트리구조이기 때문에(순환x) 최하단 자식 노드까지 내려갔다가 부모로 다시 올리면서 문제의 조건을 대입해주면 됩니다.
주의할 점은 문제의 조건상 123,456개의 노드가 존재할 수 있습니다. 그러므로 최대 재귀 깊이를 123,456이상으로 설정해주셔야 합니다. Python은 그렇습니다.
그러면 이제 순회를 돌아야겠죠?
어떻게 순회를 돌 수 있을까요?
문제의 조건대로 순회를 돌기 위한 데이터 부터 갖춰놓는게 좋겠죠.
단방향에 알맞은 부모와 자식간 연결을 나타내는 graph를 먼저 만들어야 합니다. 그리고 sheep과 wolf 정보를 알아야겠죠. (몇 번 노드에 존재하는지 알아야 합니다.) 문제 설명을 읽어보면, 두 번째 줄부터 N-1개에 줄에 2번 섬부터 N번 섬까지 섬의 정보를 나타내는라고 적혀있습니다. 따라서 각 줄이 2번부터 n번 노드까지의 정보를 갖고 있는 데이터임을 알 수 있습니다. 이를 통해 sheep과 wolf를 따로 저장하면, 더할 값과 뺄 값을 따로 관리할 수 있습니다.
결국 알고리즘을 사용하기 전 사전에 준비할 코드는 다음과 같습니다.
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline
n = int(input())
graph = [[] for _ in range(n+1)] # tree
s = [0 for _ in range(n+1)] # Sheep
w = [0 for _ in range(n+1)] # Wolf
for i in range(2, n+1): # 나머지 데이터 받기
sw, number, to = map(str, input().split())
# 문제의 조건에 따라 sw가 S일 때와 W일 때를 구분해서 따로 값을 저장합니다!
if sw == "S": s[i] = int(number)
else: w[i] = int(number)
# graph의 경우 단방향으로만 저장합니다.
# 이때 밑으로 내려가는 과정을 먼저 진행하기 때문에
# index는 parent, value는 child로 생각하고 저장해주어야 합니다.
graph[int(to)].append(i)이를 통해 DFS를 돌리면 되겠죠.
💻 DFS 코드 짜기
가장 먼저 할 것이 최하단 노드까지 내려가 주어야 합니다.
조건 1. 1번 노드(최상단 노드)에서 graph를 순회하면서 연결된 노드로 쭉쭉 내려가 줍시다!
조건 2. 최하단 노드까지 내려왔으면, 더 이상 아래로 내려갈 노드가 존재하지 않으므로 DFS를 돌지 않습니다. 그러면 이제 문제에서 주어진 상황을 대입하면서 연산을 수행하여 부모노드로 올려주면 됩니다.
조건 3. 부모 노드로 올려주다가 최상단 노드로 돌아오면(무조건 sheep, wolf 둘 다 0인 노드) 이를 answer에 더해줍시다.
이러한 방식을 코드로 옮기면 다음과 같습니다.
answer = 0
def dfs(start):
global answer
# 살아 남은 양을 기록할 number 생성
number = 0
# 조건 1. DFS를 돌 때마다 살아 남은 양을 더해주어야 함, 최하단 노드까지 내려가야 함.
for nxt in graph[start]:
number += dfs(nxt)
# 조건 2.
# 자식에서 부모로 올라올 때 해당 노드가 늑대를 갖고 있으면 늑대의 수만큼 빼주되, 0보다 작거나 같으면 0을 return 함
if w[start] != 0:
return number - w[start] if number - w[start] > 0 else 0
# 자식에서 부모로 올라올 때 해당 노드가 양을 갖고 있으면, number에서 양의 수만큼만 더해서 return 함
if s[start] != 0:
return number + s[start]
# 조건 3.
# 최상단 노드로 올라왔으면, answer에 number를 더해줌
if start == 1:
answer += number
return마지막으로 위 문제를 풀기 위한 최종 코드입니다.(python으로 돌리시길 바랍니다. pypy는 메모리 초과가 납니다.)
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline
n = int(input())
graph = [[] for _ in range(n+1)]
s = [0 for _ in range(n+1)]
w = [0 for _ in range(n+1)]
for i in range(2, n+1):
sw, number, to = map(str, input().split())
if sw == "S": s[i] = int(number)
else: w[i] = int(number)
graph[int(to)].append(i)
answer = 0
def dfs(start):
global answer
number = 0
for nxt in graph[start]:
number += dfs(nxt)
if w[start] != 0:
return number - w[start] if number - w[start] > 0 else 0
if s[start] != 0:
return number + s[start]
if start == 1:
answer += number
return
dfs(1)
print(answer)