[프로그래머스] 모두 0으로 만들기 - Java

문제 설명

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

문제

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

• 임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한쪽은 1 감소시킵니다.

하지만, 모든 트리가 위의 행동을 통하여 모든 점들의 가중치를 0으로 만들 수 있는 것은 아닙니다. 당신은 주어진 트리에 대해서 해당 사항이 가능한지 판별하고, 만약 가능하다면 최소한의 행동을 통하여 모든 점들의 가중치를 0으로 만들고자 합니다.

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

제한사항

알고리즘: DFS, Tree

풀이

트리의 모든 정점의 가중치를 0으로 만드는데 만들 수 없다면 -1, 만들 수 있다면 최소 몇 번만에 가능한지를 찾는 문제이다.

먼저 0으로 만들 수 없는 경우부터 보자,
모든 정점의 가중치의 합이 0이 아니라면 만들 수 없다. (결국 문제에서 제시한 연산은 연결된 노드중 한 노드가 증가하면 한 노드는 감소하기 때문에)

가중치의 합이 0이라면 무조건 0으로 만들 수 있는 트리이다.
이 경우는 최소 몇 번만에 가능한지를 찾아야 하는데,
최소 연산을 보장하기 위해서 우리는 자식 노드의 가중치를 다 더해가면서 루트 노드까지 올릴 것이며 연산을 했을 때 자식 노드의 가중치를 0으로 만들어 줄 것이다. -> (불필요한 연산을 줄이기 위해서)
이때 자식 노드 가중치를 0으로 만들기 위해서 필요한 연산 횟수를 answer에 더해주면 된다. (가중치의 절댓값)

그런데 여기서 루트 노드부터 탐색하려면 루트 노드를 알아야 하는데, 트리를 잘 생각해 보면 쉽게 찾아낼 수 있다. 결론만 말하자면 모든 노드는 루트 노드가 될 수 있다. 즉 어디서 시작하든 똑같은 결괏값을 가진다는 얘기다.

소스 코드

import java.util.*;
class Solution {
    static ArrayList<ArrayList<Integer>> tree = new ArrayList<>();
    static boolean[] visited;
    static long answer = 0;
    public long solution(int[] w, int[][] edges) {
        long sum = 0;
        visited = new boolean[w.length];
        for(int i=0; i<w.length; i++) {
            tree.add(new ArrayList<>());
            sum += w[i];
        }
        if(sum == 0) {
            for(int i=0; i<edges.length; i++) {
                int a = edges[i][0];
                int b = edges[i][1];
                tree.get(a).add(b);
                tree.get(b).add(a);
            }
            visited[0] = true;
            DFS(0, w);
        } else return -1;
        
        return answer;
    }
    
    static long DFS(int a, int[] w) {
        long sum_weight = w[a];
        for(int i=0; i<tree.get(a).size(); i++) {
            int b = tree.get(a).get(i);
            if(!visited[b]) {
                visited[b] = true;
                sum_weight += DFS(b, w);
            }
        }
        answer += Math.abs(sum_weight);
        return (long) sum_weight;
    }
}