[Programmers] 양과 늑대 (Java)

1.  문제 링크

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

2.  문제

2진 트리 모양 초원의 각 노드에 늑대와 양이 한 마리씩 놓여 있습니다. 이 초원의 루트 노드에서 출발하여 각 노드를 돌아다니며 양을 모으려 합니다. 각 노드를 방문할 때 마다 해당 노드에 있던 양과 늑대가 당신을 따라오게 됩니다. 이때, 늑대는 양을 잡아먹을 기회를 노리고 있으며, 당신이 모은 양의 수보다 늑대의 수가 같거나 더 많아지면 바로 모든 양을 잡아먹어 버립니다. 당신은 중간에 양이 늑대에게 잡아먹히지 않도록 하면서 최대한 많은 수의 양을 모아서 다시 루트 노드로 돌아오려 합니다.

예를 들어, 위 그림의 경우(루트 노드에는 항상 양이 있습니다) 0번 노드(루트 노드)에서 출발하면 양을 한마리 모을 수 있습니다. 다음으로 1번 노드로 이동하면 당신이 모은 양은 두 마리가 됩니다. 이때, 바로 4번 노드로 이동하면 늑대 한 마리가 당신을 따라오게 됩니다. 아직은 양 2마리, 늑대 1마리로 양이 잡아먹히지 않지만, 이후에 갈 수 있는 아직 방문하지 않은 모든 노드(2, 3, 6, 8번)에는 늑대가 있습니다. 이어서 늑대가 있는 노드로 이동한다면(예를 들어 바로 6번 노드로 이동한다면) 양 2마리, 늑대 2마리가 되어 양이 모두 잡아먹힙니다. 여기서는 0번, 1번 노드를 방문하여 양을 2마리 모은 후, 8번 노드로 이동한 후(양 2마리 늑대 1마리) 이어서 7번, 9번 노드를 방문하면 양 4마리 늑대 1마리가 됩니다. 이제 4번, 6번 노드로 이동하면 양 4마리, 늑대 3마리가 되며, 이제 5번 노드로 이동할 수 있게 됩니다. 따라서 양을 최대 5마리 모을 수 있습니다.
각 노드에 있는 양 또는 늑대에 대한 정보가 담긴 배열 info, 2진 트리의 각 노드들의 연결 관계를 담은 2차원 배열 edges가 매개변수로 주어질 때, 문제에 제시된 조건에 따라 각 노드를 방문하면서 모을 수 있는 양은 최대 몇 마리인지 return 하도록 solution 함수를 완성해주세요.

3.  제한사항

• 2 ≤ info의 길이 ≤ 17
  • info의 원소는 0 또는 1 입니다.
  • info[i]는 i번 노드에 있는 양 또는 늑대를 나타냅니다.
  • 0은 양, 1은 늑대를 의미합니다.
  • info[0]의 값은 항상 0입니다. 즉, 0번 노드(루트 노드)에는 항상 양이 있습니다.
• edges의 세로(행) 길이 = info의 길이 - 1
  • edges의 가로(열) 길이 = 2
  • edges의 각 행은 [부모 노드 번호, 자식 노드 번호] 형태로, 서로 연결된 두 노드를 나타냅니다.
  • 동일한 간선에 대한 정보가 중복해서 주어지지 않습니다.
  • 항상 하나의 이진 트리 형태로 입력이 주어지며, 잘못된 데이터가 주어지는 경우는 없습니다.
  • 0번 노드는 항상 루트 노드입니다.

입출력 예

4.  소스코드

import java.util.*;

class Solution {
    static HashMap<Integer, ArrayList<Integer>> map;
    static int maxCnt;
    public int solution(int[] info, int[][] edges) {
        map = new HashMap<>();
        for(int i = 0; i < edges.length; i++) {
            if(!map.containsKey(edges[i][0])) {
                map.put(edges[i][0], new ArrayList<Integer>());
            }
            map.get(edges[i][0]).add(edges[i][1]);
        }
        maxCnt = 0;
        ArrayList<Integer> list = new ArrayList<Integer>();
        list.add(0);
        dfs(0, 0, 0, list, info);
        return maxCnt;
    }
    
    public void dfs(int idx, int s, int w, ArrayList<Integer> list, int[] info) {
        if(info[idx] == 0)
            s++;
        else
            w++;
        if(s <= w)
            return;
        
        maxCnt = Math.max(maxCnt, s);
        ArrayList<Integer> next = new ArrayList<Integer>();
        next.addAll(list);
        if(map.containsKey(idx)) {
            next.addAll(map.get(idx));
        }
        next.remove(Integer.valueOf(idx));
        for(int n : next) {
            dfs(n, s, w, next, info);
        }
    }
}

5.  접근

1. 각 노드들의 연결 관계를 담은 2차원 배열 edges를 이용하여 트리를 완성합니다.
   • HashMap을 이용하여 2진 트리를 완성합니다.
   • edges의 각 배열을 확인하여 시작 노드를 key로 하고 해당 key의 value인 ArrayList에 도착 노드를 추가해줍니다.
2. DFS를 이용하여 경로들을 방문해보고 그 중 모을 수 있는 양은 최대 몇 마리인지 구합니다.
   • dfs 함수는 현재 노드(idx), 양의 수(s), 늑대의 수(w), 현재 노드에서 다음 탐색할 노드들을 담은 ArrayList(ArrayList list), info 정보를 인자로 받습니다.
   • 현재 노드에 양이 있는지 늑대가 있는지 확인하여 해당 동물의 수를 1 늘려줍니다.
   • 만약 늘린 후에 늑대의 수가 양의 수보다 크거나 같다면 더이상 진행할 수 없기 때문에 해당 재귀함수에 대한 실행은 끝납니다.
   • 양의 수가 늑대의 수보다 많다면 최대로 모을 수 있는 양의 수를 나타내는 변수 maxCnt의 값을 현재 maxCnt와 현재 양의 수 중 더 큰 값으로 변경합니다. maxCnt는 처음에는 0으로 초기화됩니다.
   • 현재 노드에서 다음 탐색할 노드들을 담을 ArrayList next를 생성하고 인자로 전달받은 다음 탐색할 노드들을 next에 담아줍니다.
   • 만약 현재 노드가 map에 key로 존재한다면, 즉 leaf 노드가 아니고 현재 노드에 연결되어 이동할 노드가 있다면 이동할 노드들 또한 next에 담아주고 인자로 전달된 노드들에는 현재 노드 또한 포함되어있기 때문에 현재 노드는 next에서 제거해줍니다.
   • next에 있는 각 노드들을 다음 방문할 노드로 선택하여 해당 노드를 현재 노드로 하는 재귀함수를 호출합니다.
   • 재귀함수 호출을 통해 전달받은 현재 노드에서 위 작업들을 수행하며 늑대의 수가 양의 수보다 같거나 많아질 때까지 각 경로들을 진행해보고 재귀함수 호출이 모두 종료됐을 때 maxCnt 값을 확인합니다.