모두 0으로 만들기

import java.util.*;
class Solution {
    private static boolean[] visited;
    private static ArrayList<Integer>[] edgeList;
    private static long[] L_a;
    private static long answer = 0;
 
    public static long DFS(int index) {
        visited[index] = true;
        // 현재 노드와 인접한 노드를 탐색하며 가중치 갱신
        for (int i = 0; i < edgeList[index].size(); i++) {
            int adjNode = edgeList[index].get(i);
            // 방문하지 않은 경우만 수행
            if (visited[adjNode] != true) {
                L_a[index] += DFS(adjNode);
            }
        }
        answer += Math.abs(L_a[index]);
        return L_a[index];
    }
    
    public long solution(int[] a, int[][] edges) {
        int len = a.length;
        edgeList = new ArrayList[len];
        visited = new boolean[len];
        L_a = Arrays.stream(a).asLongStream().toArray();
 
        long sum = 0;
        // sum 체크 및 노드 추가
        for (int i = 0; i < len; i++) {
            sum += L_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;
    }
}

연결된 노드의 가중치를 모두 계산하여 0으로 만드는 문제이다.

한 번 접근했던 노드를 visited 변수를 통해 바로 체크하였고, 계산 과정에서 범위 초과를 대비하여 long으로 변환하였다.

 

visited 변수를 사용하지 않아서 시간 초과가 난 코드는 아래와 같다.

import java.util.*;
class Solution {
    public long solution(int[] a, int[][] edges) {
        long answer = 0;
        int len = a.length;
        ArrayList<Integer>[] edgeList = new ArrayList[len];
        ArrayList<Integer> leafs = new ArrayList<Integer>();
        long[] L_a = Arrays.stream(a).asLongStream().toArray();
 
        long sum = 0;
        // sum 체크 및 노드 추가
        for (int i = 0; i < len; i++) {
            sum += L_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]);
        }
 
        for (int i = 0; i < len; i++) {
            // leaf 노드일 때
            if (edgeList[i].size() == 1)
                leafs.add(i);
        }
 
        // 마지막 노드가 남을 때 까지 수행
        while (leafs.size() != 1) {
            int children = leafs.remove(0);
            int parent = edgeList[children].get(0);
 
            // 가중치 갱신
            L_a[parent] += L_a[children];
            answer += Math.abs(L_a[children]);
            // 자식 노드 제거
            edgeList[parent].remove((Object) children);
 
            // 부모가 leaf 노드가 됐을 경우 추가
            if (edgeList[parent].size() == 1)
                leafs.add(parent);
        }
        return answer;
    }
}

노드를 추가하고 제거하는데에 많은 시간이 소요된 것으로 보인다.

 

참고 사이트 : https://javariankm.tistory.com/41