[프로그래머스 Level3] 양과 늑대 - 파이썬(Python) - 2022 KAKAO BLIND RECRUITMENT

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

프로그래머스

문제 설명

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 함수를 완성해 주세요.

 

 

 

---

처음에는 정말 복잡하게 생각했던 문제였다. bfs는 안되지 않을까 했는데 사실 bfs/dfs 모두 가능한 문제였다.

나는 dfs가 더 적합하다고 생각해서 dfs로 풀이했다.

 

dfs에는 현재 노드 위치와 양의 수, 늑대의 수, 이동할 수 있는 모든 노드들을 인자로 담았다. 늑대의 수가 양의 수의 이상이 되면 모든 양들이 잡아먹히므로 return 했고, 가능한 양의 수는 매번 갱신했다.

 

 

어려웠던 부분은 이 부분이다. 이동가능한 노드들에서 이제 이동할 장소를 뺀 것들의 집합과 graph에서 이제 이동할 장소의 자식노드들의 집합을 합집합한 위치들이 다음에 이동할 수 있는 장소들이다.

dfs(pos, sheep+1, wolf, set(tree[pos])|can_move-{pos})
dfs(pos, sheep, wolf+1, set(tree[pos])|can_move-{pos})

 

 

 

 

<< 최종 코드 >>

from collections import defaultdict

answer = 0

def solution(info, edges):
    tree = defaultdict(list)
    for parent, child in edges:
        tree[parent].append(child)
    
    def dfs(cur, sheep, wolf, can_move):
        global answer
        if wolf >= sheep:
            return
        answer = max(answer, sheep)
        
        for pos in can_move:
            if info[pos] == 0:
                dfs(pos, sheep+1, wolf, set(tree[pos])|can_move-{pos})
            else:
                dfs(pos, sheep, wolf+1, set(tree[pos])|can_move-{pos})
                        
    dfs(0, 1, 0, set(tree[0]))
    return answer