[트리] 모두 0으로 만들기 - 월간 코드 챌린지 시즌 2

모두 0으로 만들기

문제 소개

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다.

임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다.
하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.

트리의 각 점의 가중치를 의미하는 1차원 정수 배열 a와 트리의 간선 정보를 의미하는 edges가 매개변수로 주어집니다. 주어진 행동을 통해 트리의 모든 점들의 가중치를 0으로 만드는 것이 불가능하다면 -1을, 가능하다면 최소 몇 번만에 가능한지를 찾아 return 하도록 solution 함수를 완성해주세요. (만약 처음부터 트리의 모든 정점의 가중치가 0이라면, 0을 return 해야 합니다.)

시나리오

문제를 막 접했을 때 연결된 노드의 수가 가장 작은 노드부터 가중치를 0으로 바꿔가겠다는 아이디어는 떠올렸다. 하지만 최소 노드를 반복해서 탐색했기 때문에 결과는 시간 초과였다. 이후, 트리 구조를 상기시키고 dfs에 적용하여 풀이했다.

1. 가중치의 총 합이 0이 아니라면 불가능하다. -1 반환.
2. 연결된 노드의 수가 가장 작은 노드부터 가중치를 0으로 바꿔간다.
   • 즉, 트리 구조에서 단말 노드부터 루트 노드까지 올라가며 가중치를 바꿔간다. 후위 순회(post-order)와 같다.
   • 어떤 노드를 root로 설정해도 트리 구조가 유지된다.

문제 풀이

1. 재귀

import sys
sys.setrecursionlimit(300000)

answer = 0

def dfs(node, a, tree, visited):
    global answer
    visited[node] = 1
    
    for leaf in tree[node]:
        if visited[leaf]: continue
        a[node] += dfs(leaf, a, tree, visited)
    answer += abs(a[node])
        
    return a[node]
    

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

    global answer
    tree = [set() for _ in range(len(a))]
    visited = [0 for _ in range(len(a))]
    
    for u, v in edges:
        tree[u].add(v)
        tree[v].add(u)
        
    dfs(0, a, tree, visited)
    
    return answer

2. 반복문

def solution(a, edges):
    answer, root = 0, 0
    stack = [root]
    tree = [set() for _ in range(len(a))]
    visited = [False for _ in range(len(a))]
    
    if sum(a) != 0: return -1

    # 트리 구조
    for u, v in edges:
        tree[u].add(v)
        tree[v].add(u)

    while stack:
        node = stack.pop()

        if visited[node] and len(tree[node]) == 1:
            p_node = tree[node].pop()
            tree[p_node].discard(node)      # 연결 간선 삭제
            a[p_node] += a[node]
            answer += abs(a[node])
            a[node] = 0
        else:
            visited[node] = True
            if node != root:                
                stack.append(node)
            for l_node in tree[node]:
                if not visited[l_node]:
                    stack.append(l_node)
    return answer

회고

재귀로 풀 수 있는 모든 문제는 반복문으로도 풀 수 있다고 한다. 이번 문제는 복잡해서 그런지 반복문 풀이를 쉽게 찾아볼 수 없었다. 직접 변환해보면서 머리를 싸맸지만...덕분에 스택을 사용하는 dfs를 연습해봤다. 결국 성공. 도전하길 잘했다. 😐