Programmers - 모두 0으로 만들기

SJ0000·2022년 6월 26일

문제 링크

나는 dfs로 풀었는데 위상정렬로도 풀 수 있다고 한다.
dfs로 탐색 하다가 더 이상 갈 곳이 없을 때 자신의 가중치를 parent에 전달한다.
parent에 전달하면서 total을 업데이트한다.

import java.util.*;

class Solution {

    List<Node> allNodes = new ArrayList<>();
    List<Node>[] edgeInfo;
    static boolean[] visit;
    long total = 0;
    public long solution(int[] a, int[][] edges) {
        if(!isPossible(a))
            return -1;
        init(a, edges);
        dfs(0, -1);
        return total;
    }

    private void dfs(int node, int parent){
        visit[node] = true;

        List<Node> nextNodes = edgeInfo[node];
        for(int i=0;i<nextNodes.size();i++){
            Node next = nextNodes.get(i);
            if(!visit[next.num])
                dfs(next.num, node);
        }
        if(parent==-1)
            return;

        Node p = allNodes.get(parent);
        Node n = allNodes.get(node);
        p.weight += n.weight;
        total += Math.abs(n.weight);
    }

    private boolean isPossible(int[] a){
        long sum = 0;
        for(int i=0;i<a.length;i++){
            sum += a[i];
        }
        return sum == 0;
    }

    private void init(int[] a, int[][] edges) {
        edgeInfo = new ArrayList[a.length];
        visit = new boolean[a.length];

        for(int i=0;i<edgeInfo.length;i++){
            edgeInfo[i] = new ArrayList<>();
        }

        for (int i = 0; i < a.length; i++) {
            allNodes.add(new Node(i, a[i]));
        }
        for (int[] edge : edges) {
            Node x = allNodes.get(edge[0]);
            Node y = allNodes.get(edge[1]);
            edgeInfo[x.num].add(y);
            edgeInfo[y.num].add(x);
        }
    }

    static class Node {
        int num;
        long weight;
        public Node(int num, long weight) {
            this.num = num;
            this.weight = weight;
        }
    }
}
profile
잘하고싶은사람

0개의 댓글