모두 0으로 만들기(프로그래머스) JAVA

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

프로그래머스

 

문제 설명

 

문제에 대한 아이디어

위의 그림과 다르게 0이 root인 트리로 만들어봤다.

위 문제를 보면 어떤 노드를 루트로 선택하든 문제가 없어 보인다. 나는 가장 쉽게 노드 0을 선택해서 위의 그림처럼 트리를 만들었다. 고민을 하다가 생각해 낸 방법은 "모든 노드를 0으로 만들려면 리프 노드부터 0으로 만들면 되지 않을까" 였다.

• 리프 노드의 값을 parent로 보내 버린다.
• 즉 (2,2)와 (4,2)의 각각 값을 parent인 (3,1)로 보내버린다. 
• 그러면 리프 노드 부터 모두 0으로 바뀌고 (3,5)가 된다. 
• 위를 반복하면 결국 노드 0까지 가게된다.
• 노드 0에 도착했을 때, 노드 0이 0이면 true, 0이 아니라면 false를 리턴한다.

 

구현

import java.util.*;
class Node{
    int num;
    long count;
    public Node(int num, long count){
        this.num = num;
        this.count = count;
    }
}

class Solution {
    ArrayList<Integer> pro[]; // 트리
    ArrayList<Node> arr = new ArrayList<>(); // index에 맞는 정보 저장 (pair라고 생각하면 됨)
    boolean visited[];
    long answer = 0;
    public long solution(int[] a, int[][] edges) {
        pro = new ArrayList[a.length];
        visited = new boolean[a.length];
        //트리 초기화
        for(int i =0; i< a.length; ++i){
            pro[i] = new ArrayList<>();
        }
        // 각 index에 알맞는 값 세팅
        for(int i = 0; i<a.length; ++i){
            arr.add(new Node(i,a[i]));
        }
        //트리 만들기
        for(int i = 0; i< edges.length; ++i){
            int u = edges[i][0];
            int v = edges[i][1];
            pro[u].add(v);
            pro[v].add(u);
        }
        if(!DFS(0,0)){
            return -1;
        }
        else return answer;
    }
    
    public boolean DFS(int now, int parent){
        visited[now] = true;
        for(int i = 0; i< pro[now].size(); ++i){
            int next = pro[now].get(i);
            if(visited[next]) continue;
            DFS(next,now);
        }
        // 이 부분을 실행 할 때는 더 이상 갈 곳이 없다는 것이므로 리프노드이다.
        if(now == parent){ // 노드 0 에 도착했을 때
            if(arr.get(now).count != 0) return false;
        }
        // 노드 0이 아니라면
        arr.get(parent).count += arr.get(now).count;
        answer += Math.abs(arr.get(now).count); // 음수도 존재하기 때문에 절댓값을 더해야 한다.
        arr.get(now).count = 0;
        return true;
    }
}