[알고리즘] 프로그래머스 3단계 - 양과 늑대

프로그래머스 문제 링크

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

 

문제 설명

양과 늑대가 이진 트리 형식의 배열에 존재한다.

늑대가 양의 수와 같거나 많은 경우 이동할 수 없고 양의 수가 많아야 한다.

양의 수가 최대인 경우 몇마리인지 리턴한다.

초기 루트 노드는 양으로 시작한다.

info에는 양인지 늑대인지 여부가 0과 1로 구성되어있고 edges에는 각 노드간에 연결된 정보와 info에 해당되는 양인지 늑대인지 여부가 결정된다.

 

예시

 

info
[0,0,1,1,1,0,1,0,1,0,1,1]

edges
[[0,1],[1,2],[1,4],[0,8],[8,7],[9,10],[9,11],[4,3],[6,5],[4,6],[8,9]]

result
5

import java.util.*;

public class Main {

    private static ArrayList<Integer>[] tree;

    public static void main(String[] args) {
        System.out.println(solution(new int[]{0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0},new int[][]{{0,1},{0,2},{1,3},{1,4},{2,5},{2,6},{3,7},{4,8},{6,9},{9,10}}));
    }

    public static int solution(int[] info, int[][] edges) {
        buildTree(info, edges);

        int answer = 0;

        ArrayDeque<Info> queue = new ArrayDeque<>();
        queue.add(new Info(0, 1, 0, new HashSet<>()));

        while (!queue.isEmpty()) {
            Info now = queue.poll();

            answer = Math.max(answer, now.sheep);

            // 맨처음 노드가 0번일때는 형제노드가 없어서 아래 addAll 함수 호출시 하위 노드만 포함되게 되지만
            // 한번 하위 노드를 돌고 형제 노드가 있다면 포함 시킨다.
            // 즉 이전에 있던 형제 노드와 현재 노드에 연결된 하위 노드 합쳐지는 배열이다.
            now.visited.addAll(tree[now.node]);

            // 형제 노드 + 하위 노드를 하나씩 꺼내서 늑대인지 양인지 판별 후 수를 비교 하여 조건에 따라서 큐에 담는다.
            for (int next : now.visited) {
                // 여기가 중요한 부분인데 현재 노드의 정보만 제거 하고 연결된 모든 노드 정보를 새로 생성될 큐에 추가한다.
                HashSet<Integer> set = new HashSet<>(now.visited);
                set.remove(next);

                // 양인지 늑대인지 판별
                // 책에서는 중첩 if문으로 구성되었지만 문제에서 기본적으로 0아니면 1만 확실히 주기 때문에 다음과 같이 연산횟수를 줄일 수 있다.
                //if (info[next] == 1) {
                //    if (now.sheep != now.wolf + 1) {
                //        queue.add(new Info(next, now.sheep, now.wolf + 1, set));
                //    }
                //} else {
                //    queue.add(new Info(next, now.sheep + 1, now.wolf, set));
                //}
                
                if (info[next] == 0) {
                    queue.add(new Info(next, now.sheep + 1, now.wolf, set));
                    // 양보다 숫자가 적어야 늑대를 큐에 담는다.
                } else if(now.sheep != now.wolf + 1) {
                    queue.add(new Info(next, now.sheep, now.wolf + 1, set));
                }
            }
        }
        return answer;
    }

    private static void buildTree(int[] info, int[][] edges) {
        tree = new ArrayList[info.length];
        // 각 노드에 연결된 하위 노드 목록 넣기 위한 리스트 초기화 작업
        for (int i = 0; i < tree.length; i++) {
            tree[i] = new ArrayList<>();
        }

        // 각 노드 별로 연결된 노드 추가 ( 0 노드 부터 하위 노드 추가 하여 연결)
        for (int[] edge : edges) {
            tree[edge[0]].add(edge[1]);
        }
    }

    private static class Info {
        int node, sheep, wolf;
        HashSet<Integer> visited;

        public Info(int node, int sheep, int wolf, HashSet<Integer> visited) {
            this.node = node;
            this.sheep = sheep;
            this.wolf = wolf;
            this.visited = visited;
        }
    }
}

 

풀다가 어려워서 책 풀이를 확인했다. 풀이를 확인해도 foreach문 내부에서 HashSet설정하는 이유를 한참 헤맸는데 트리에 익숙해질 필요가 있겠다.