[프로그래머스/파이썬(Python)] 모두 0으로 만들기

https://programmers.co.kr/learn/courses/30/lessons/76503

코딩테스트 연습 - 모두 0으로 만들기

 

트리구조와 DFS를 사용하는 문제입니다. 최근에 기업 코딩테스트에서 트리 문제가 자주 나오기 때문에 잘 알아두면 좋을 것 같습니다.

 

이 문제의 특징은 1) 어떤 정점을 잡아도 해당 정점이 루트 노드가 될 수 있다는 점, 2) 따라서 아무 정점이나 잡은 다음에 DFS를 통해 Leaf 노드까지 내려갈 수 있고, 3) Leaf 노드부터 가중치를 부모 노드에게 전달하면 된다는 것입니다.

 

따라서, 최초에 dfs함수에는 0번 정점을 Root라고 가정하고 가중치를 전달할 수 있습니다. 그리고 반복적으로 재귀를 수행하면서 Leaf까지 내려가고 Leaf 노드를 Call한 Parent 노드에게 가중치를 전달하면서 정답값에 더해주면 됩니다.

 

 

[풀이]

from collections import defaultdict
import sys
sys.setrecursionlimit(100000000)

answer = 0

def dfs(tree, a, visited, node):
    global answer
    visited[node] = True

    for child in tree[node]:
        if not visited[child]:
            # dfs로 리프노드까지 탐색
            leafWeight = dfs(tree, a, visited, child)

            a[node] += leafWeight
            answer += abs(leafWeight)

    return a[node]

def solution(a, edges):
    if sum(a) != 0:
        return -1

    # 정점과 연결된 edge의 개수가 1개인 것부터 그리디하게 수행
    n = len(a)      # 정점의 개수

    # 연결 정보를 담는 트리
    tree = defaultdict(list)
    for edge in edges:
        sour, dest = edge
        tree[sour].append(dest)
        tree[dest].append(sour)

    visited = [False] * n    
    dfs(tree, a, visited, 0)

    return answer