import sys
sys.setrecursionlimit(10**6)
from collections import defaultdict
def solution(a, edges):
global answer, graph
graph = defaultdict(list)
answer = 0
for v1, v2 in edges:
graph[v1].append(v2)
graph[v2].append(v1)
dfs(a, 0, -1)
return -1 if sum(a) else answer
def dfs(a, node, parent) :
global answer
for child in graph[node] :
if child == parent : continue
dfs(a, child, node)
a[parent] += a[node]
answer += abs(a[node])
a[node] = 0
풀이
트리 DFS 기본 문제다.
자식노드의 가중치를 부모노드 가중치에 더해준다.
자식노드의 가중치의 절대값을 count한다.
모든 traverse가 끝난 뒤,
return -1 if root not 0 else answer
을 해주면 된다.
import sys
sys.setrecursionlimit(10**6)
from collections import defaultdict
def solution(a, edges):
global answer, graph
graph = defaultdict(list)
answer = 0
for v1, v2 in edges:
graph[v1].append(v2)
graph[v2].append(v1)
dfs(a, 0, -1)
# sum(a) : root 가중치가 0 인지 확인
return -1 if sum(a) else answer
def dfs(a, node, parent) :
global answer
# recursive DFS part
for child in graph[node] :
if child == parent : continue
dfs(a, child, node)
# 부모노드 가중치치에 자식노드 가중치 더하기
a[parent] += a[node]
# 자식노드 가중치 절대값 count
answer += abs(a[node])
# 자식노드 가중치 0
a[node] = 0