알고리즘/PS - 프로그래머스

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

excited-hyun 2021. 8. 23. 19:45
반응형

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

 

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

각 점에 가중치가 부여된 트리가 주어집니다. 당신은 다음 연산을 통하여, 이 트리의 모든 점들의 가중치를 0으로 만들고자 합니다. 임의의 연결된 두 점을 골라서 한쪽은 1 증가시키고, 다른 한

programmers.co.kr

 

풀이

 

0번째 노드를 root로 잡고 DFS를 이용해서 leaf 노드까지 내려가서 거기서부터 0으로 만들면서 점점 root로 올라오면 되는 문제입니다.

어떤 알고리즘을 선택해야할지 고민하는 과정이 어려웠습니다ㅠㅠ

 

우선 중요한 것은 모든 가중치를 더했을 때 합이 0이 아닌 경우엔 절대 0으로 만들 수 없다는 것입니다.

그 경우를 처리해준 뒤, 연결된 edge의 경로들을 ArrayList로 연결해줍니다.

그런 뒤 dfs를 재귀로 호출하면서 leaf 노드부터 0으로 만들고 그에 따라 연결된 노드의 가중치를 업데이트해가면서

모든 노드를 0으로 만들어주면 됩니다.

 

import java.util.*;

class Solution {
    long answer = 0;
    long tmp[]; 
    int visit[];
    List<Integer>[] edgeList;
    
    public long solution(int[] a, int[][] edges) {
        
        int len = a.length;
        tmp = new long[len];
        visit = new int[len];
        long sum = 0;
        edgeList = new ArrayList[len];
        for(int i = 0; i < len; i++) {
			sum += a[i];
            tmp[i] = a[i];
            edgeList[i] = new ArrayList<Integer>();
	    }
        //불가능
        if(sum!=0)
            return -1;
        
        for(int i = 0; i < edges.length; i++) {
	    	edgeList[edges[i][0]].add(edges[i][1]);
	    	edgeList[edges[i][1]].add(edges[i][0]);
		}
        
        dfs(0);
        
        return answer;
    }
    
    long dfs(int node){
        visit[node] = 1;
        for(int i=0; i<edgeList[node].size(); i++){
            int next = edgeList[node].get(i);
            if(visit[next]==1)
                continue;
            tmp[node] += dfs(next);
            
        }
        long num = tmp[node];
        answer += Math.abs(num);
        return num;
    }
}
728x90
반응형