[프로그래머스 Level.3] 모두 0으로 만들기 (월간 코드 챌린지 시즌2) (Java)

문제 링크

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

프로그래머스

 

 

 

 

코딩테스트 연습 > 월간 코드 챌린지 시즌2 > 모두 0으로 만들기

 

 

 

문제 설명

 

 

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

 

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

 

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

 

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

 

---

제한사항

• a의 길이는 2 이상 300,000 이하입니다.
  • a의 모든 수는 각각 -1,000,000 이상 1,000,000 이하입니다.
  • a[i]는 i번 정점의 가중치를 의미합니다.
• edges의 행의 개수는 (a의 길이 - 1)입니다.
  • edges의 각 행은 [u, v] 2개의 정수로 이루어져 있으며, 이는 u번 정점과 v번 정점이 간선으로 연결되어 있음을 의미합니다.
  • edges가 나타내는 그래프는 항상 트리로 주어집니다.

 

---

 

입출력 예

a	edges	result
[-5,0,2,1,2]	[[0,1],[3,4],[2,3],[0,3]]	9
[0,1,0]	[[0,1],[1,2]]	-1

 

---

 

입출력 예 설명

 

 

입출력 예 #1

 

• 다음 그림은 주어진 트리의 모든 정점의 가중치를 0으로 만드는 과정을 나타낸 것입니다.

 

 

1. 2번 정점과 3번 정점을 선택하여 2번 정점은 1 감소시키고, 3번 정점은 1 증가시킵니다. (2번 반복)
2. 3번 정점과 4번 정점을 선택하여 4번 정점은 1 감소시키고, 3번 정점은 1 증가시킵니다. (2번 반복)
3. 0번 정점과 3번 정점을 선택하여 3번 정점은 1 감소시키고, 0번 정점은 1 증가시킵니다. (5번 반복)

 

• 모든 정점의 가중치를 0으로 만드는 데 필요한 최소 행동 횟수는 9번이므로, 9를 return 해야 합니다.

 

 

 

입출력 예 #2

 

• 주어진 트리는 모든 정점의 가중치를 0으로 만드는 것이 불가능하므로, -1을 return 해야 합니다.

 

 

 

 

 

 

나의 코드

import java.util.*;

class Solution {
    long answer;
    long[] longA;
    boolean[] visited;
    ArrayList<Integer>[] graph;
    public long solution(int[] a, int[][] edges) {
        answer = 0;
        int total = 0;
        longA = new long[a.length];
        visited = new boolean[a.length];
        graph = new ArrayList[a.length];
        
        for(int i=0; i<a.length; i++) {
            total += a[i];
            longA[i] = a[i];
            graph[i] = new ArrayList<>();
        }
        
        if(total != 0) return -1;
        
        for(int i=0; i<edges.length; i++) {
            graph[edges[i][0]].add(edges[i][1]);
            graph[edges[i][1]].add(edges[i][0]);
        }
        
        dfs(0);
        
        return answer;
    }
    
    private long dfs(int vertex) {
        visited[vertex] = true;
        
        for(int i=0; i<graph[vertex].size(); i++) {
            int next = graph[vertex].get(i);
            if(!visited[next]) longA[vertex] += dfs(next);
        }
        
        answer += Math.abs(longA[vertex]);
        
        return longA[vertex];
    }
}

 

풀이

1. 임의의 연결된 두 점을 골라 한쪽은 1 감소시키고 다른 한쪽은 1 증가시키는 것은 한쪽의 값을 다른 쪽으로 옮기는 것과 같다. 따라서, 트리 구조에서 볼 때 최종적으로 몇 번 필요한지 확인하기 위해서는 완전탐색 dfs를 이용하여 리프 노드까지 내려갔다가 올라오면서 자식 노드의 값을 부모 노드에게 옮기고 계산하며 root 노드까지 올라오고 값을 옮길 때마다 그 값을 절댓값으로 하여 카운팅해주면 된다.
2. 우선 한쪽은 -1, 한쪽은 1씩 더해주어 모두 0이 되려면 a 배열의 값을 모두 더했을 때 0이 되어야 하기 때문에 a 배열의 값을 모두 더해 0이 되지 않는다면 -1을 반환한다.
3. 자식노드의 값을 부모노드에게 더해주면서 올라오기 때문에 수가 충분히 커질 수 있기때문에 longA 배열을 만들어 a 배열의 값을 담아준다.
4. 트리 구조의 연결된 간선을 파악하기 위해서 해당 노드의 번호를 인덱스로 하여 연결된 노드의 번호들을 담기 위해 ArrayList를 원소로 갖는 배열 grpah를 만들고 edge를 순회하면서 연결된 노드를 담아준다.
5. 또한, 방문한 노드를 체크하기 위해 boolean형 배열 visited를 생성한다.
6. 이제 완전탐색 dfs를 진행하면서 리프노드까지 내려가서 자식노드의 값을 부모노드에게 옮기며 해당 값만큼 카운팅한다.
7. 루트 노드인 0노드부터 시작하여 해당 노드의 visited를 true로 변경하고 graph를 이용하여 연결된 노드들을 순회하면서 해당 연결된 노드를 방문하지 않았다면 dfs()를 재귀 호출하여 반환 받은 값을 해당 노드의 값에 더한다.
8. 재귀 호출을 하며 자식노드를 타고 리프노드까지 내려가서 해당 노드 값의 절댓값 만큼을 answer에 증가시키고 해당 노드의 값을 반환하여 재귀 호출했던 부모노드에게 전달한다.
9. 이렇게 dfs를 진행하고 확인해보면 최종적으로 모든 노드의 값을 더한 루트 노드의 값 longA[0]을 확인해보면 0이 되어 있고 값을 옮긴 수만큼(더하고 뺀 횟수) 계속해서 answer에 더해주었으므로 answer를 반환하면 된다.