[프로그래머스] 양과 늑대

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

프로그래머스

 

문제

2진 트리에서 각 노드에 늑대(1)와 양(0)이 한 마리씩 위치하고 있으며, 각 노드를 방문할 때마다 해당 노드에 있던 양과 늑대를 데리고 다니게 된다. 만약 모은 늑대의 수가 양의 수 이상이 되면 늑대가 모든 양을 잡아먹는다.
이때 루트 노드(0, 무조건 양)에서 출발하여 각 노드를 돌아다니면서 최대한 양을 많이 모을 수 있는 수를 구한다.

조건

1. info[i]는 i번 노드에 있는 양 또는 늑대 (0: 양, 1: 늑대)

2. info[0]은 항상 0

3. int[][] edges = [[부모 노드 번호, 자식 노드 번호], ...]

4. 0번 노드가 항상 루트 노드

 

생각 정리해보기

[기본 규칙] 양이 늑대보다 많은 상태를 유지할 수 있을 경우에만 이동할 수 있다.

그러니까 다른 노드로 이동하기 전에 "양이 늑대보다 많은지"를 확인해야 한다.

 

[이어서] 지금 당장 방문하지 못하는 노드가 있다.

지금까지 BFS, DFS 탐색에서는 왼쪽 > 오른쪽으로 노드를 방문해왔는데, 이번 경우에는 그 방향으로 진행했을 땐 방문하지 못했지만 나중에는 방문할 수 있는 경우가 존재한다.

예를 들면 현재 상태: 양 2 & 늑대 1일 경우,

- 왼쪽 자식 노드에 늑대: 당장 방문 불가능

- 오른쪽 자식 노드에 양: 오른쪽 자식 노드에 방문하면 양 3 & 늑대 1이 되므로 이후 왼쪽 자식 노드에 방문할 수 있음

그러니까 당장 방문하지 못하더라도 영원히 방문하지 못하는 것이 아니라 나중으로 "보류"하게 되므로 후보로 남겨두어야 한다.

 

[이어서] 아직 방문하지 못하는 노드를 목록으로 관리한다.

"아직 방문하지 못하는"이라고 명시하긴 했으나 사실 다음 차례에 방문을 시도해 볼 노드 목록이라고 보면 된다.

BFS에서 다음 방문 예정 노드를 큐에 등록해서 관리했던 것과 비슷하다. 그러나 방문하지 못한 노드는 큐에 계속 남겨두어야 한다.

따라서 "늑대 수 때문에 다음에 다시 시도해 볼 노드" + "현재 방문한 노드의 연결 노드"로 계속 업데이트하게 된다.

 

[이어서] 목록에 있는 노드를 순서대로 방문한다.

방문 전에는 양과 늑대 상태를 확인하고, 방문한 노드는 목록에서 삭제한다.

 

결국 시도해 볼 모든 노드들을 후보 목록으로 들고 다니면서 + 지금 못 가면 나중에 다시 시도하는 방식이 포인트라고 생각한다.

 

 

DFS 구현

양, 늑대 상태를 확인하고 돌아오는 방식을 적용하기 간편한 DFS로 구현했다.

 

1. 먼저 부모 노드와 자식 노드의 엣지만 주어진 목록을 연결 리스트 기반 그래프로 변환한다.

Map<Integer, List<Integer>> graph = new HashMap<>();
for(int[] edge:edges) {
	graph.computeIfAbsent(edge[0], k->new ArrayList<>()).add(edge[1]);
}

 

2. DFS에서는 방문 시도할 노드 목록과 양과 늑대 수(상태)를 들고 다니도록 구현한다.

- 후보 목록의 노드를 차례대로 탐색한다. (방문할 수 있는지 확인)

- 양<=늑대라면 보류한다.

- 양>늑대라면 현재 후보 목록에서 해당 노드를 제거하고 (방문할 거니까), 다시 해당 노드의 자식 노드를 추가한 새로운 후보 목록을 추가하여 방문한다.

- 방문한 노드에서는 현재 양 수를 저장된 최대 마릿수와 비교하고 크다면 업데이트한다.

void dfs(int[] info, List<Integer> q, int s, int w) {
    answer = Math.max(answer, s); //최댓값으로 양 수 업데이트
	
    //현재 후보 목록 차례대로 방문 시도 
    for(int i=0; i<q.size(); i++) {
        int newSheep = s; int newWolf = w;
		
        int next = q.get(i);
        if(info[next] == 1) newWolf++; else newSheep++; //다음 노드에 있는 동물 연산

        if(newSheep <= newWolf) continue; //양<=늑대면 방문 보류

		//양>늑대이면 방문!
        List<Integer> newQ = new ArrayList<>(q); 
        newQ.remove(i); //현재 후보 목록에서 다음 노드 제거
        if(graph.containsKey(next)) newQ.addAll(graph.get(next)); //다음 노드의 자식 노드 추가
 
        dfs(info, newQ, newSheep, newWolf); //방문
}

 

 

전체 코드

class Solution {
    Map<Integer, List<Integer>> graph = new HashMap<>();
    int answer = 0;
    
    public int solution(int[] info, int[][] edges) {
        //그래프
        for(int[] edge:edges) graph.computeIfAbsent(edge[0], k->new ArrayList<>()).add(edge[1]);
        
        dfs(info, graph.get(0), 1, 0);
        
        return answer;
    }
    
    void dfs(int[] info, List<Integer> q, int s, int w) {
        answer = Math.max(answer, s);
        
        for(int i=0; i<q.size(); i++) {
            int newSheep = s; int newWolf = w;
            
            int next = q.get(i);
            if(info[next] == 1) newWolf++; else newSheep++;
            
            if(newSheep <= newWolf) continue;
            
            List<Integer> newQ = new ArrayList<>(q);
            newQ.remove(i);
            if(graph.containsKey(next)) newQ.addAll(graph.get(next));
            
            dfs(info, newQ, newSheep, newWolf);
        }
    }
    
}