[Python | Level3] 양과 늑대 / 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 함수를 완성해주세요.

---

제한사항

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

---

문제 풀이

"양과 늑대" 문제는 트리 구조를 기반으로 하는 완전탐색 문제이다. 각 노드는 양 또는 늑대를 나타내며, 루트 노드인 0번 노드에서 출발하여 양을 최대한 많이 수집하는 것이 목표이다. 단, 탐색 도중에 늑대의 수가 양의 수와 같아지거나 많아지면 더 이상 진행할 수 없다는 제약이 주어진다. 따라서 단순한 트리 탐색만으로는 해결할 수 없고, 탐색 중 양과 늑대의 수를 상태로 관리하면서 가능한 모든 경우를 탐색해야 한다.

문제에서 제공되는 정보는 두 가지다. 하나는 info 배열로, 각 노드가 양인지 늑대인지를 알려준다. 다른 하나는 edges 배열로, 부모-자식 관계를 나타낸다. 이를 기반으로 트리를 구성한 뒤, 루트 노드부터 가능한 경로를 따라가며 DFS(깊이 우선 탐색)를 수행한다. 이때 핵심은 현재 탐색 경로에서 갈 수 있는 후보 노드들을 항상 기억하고, 그중 하나를 선택해 탐색을 이어가는 것이다. 단순히 현재 노드의 자식 노드로만 이동하는 것이 아니라, 이전에 방문했던 경로에서 파생된 모든 자식 노드 중에서 선택할 수 있어야 하므로, 현재까지 방문 가능한 후보 노드 리스트를 별도로 유지해야 한다.

DFS 함수는 다음과 같은 상태를 인자로 받는다: 현재 위치한 노드 번호, 지금까지 수집한 양의 수, 늑대의 수, 그리고 이후에 방문 가능한 노드들의 리스트이다. 함수가 호출되면 먼저 현재 노드가 양인지 늑대인지 판별하여 각 수를 갱신한다. 이때 늑대의 수가 양의 수 이상이 되면 해당 경로는 유효하지 않으므로 탐색을 중단한다. 이후 현재 노드의 자식 노드들을 후보 리스트에 추가하고, 현재 노드를 후보 리스트에서 제거한다. 후보 노드 리스트는 앞으로 방문 가능한 노드들을 의미하며, 여기서 하나씩 꺼내 재귀적으로 DFS를 호출함으로써 모든 가능 경로를 탐색할 수 있다. 이때 각각의 재귀 호출에서는 후보 노드 리스트를 복사하여 전달해야 서로 다른 탐색 경로가 영향을 주지 않도록 한다.

이 과정을 반복하면 가능한 모든 유효 경로를 탐색할 수 있고, 그중 양을 가장 많이 수집할 수 있었던 경우의 양 수를 정답으로 반환하면 된다. 따라서 DFS를 수행하면서 매번 최대 양 수를 갱신하는 방식으로 구현한다. 이 문제는 단순한 방문 체크나 경로 저장이 아닌 상태 기반 완전탐색이 핵심이며, 트리 구조의 비정형적인 탐색을 요하므로 DFS 호출 시 후보 노드 리스트를 관리하는 구현의 정확도가 중요하다.

---

구현

def solution(info, edges):
    from collections import defaultdict

    # 트리 구조 저장
    graph = defaultdict(list)
    for parent, child in edges:
        graph[parent].append(child)

    max_sheep = 0

    def dfs(current_node, sheep, wolf, next_nodes):
        nonlocal max_sheep

        # 현재 노드가 양인지 늑대인지 판단
        if info[current_node] == 0:
            sheep += 1
        else:
            wolf += 1

        # 양보다 늑대가 같거나 많아지면 종료
        if wolf >= sheep:
            return

        # 최대 양 수 갱신
        max_sheep = max(max_sheep, sheep)

        # 현재 노드에서 갈 수 있는 자식 노드 추가
        candidates = next_nodes + graph[current_node]
        # 현재 노드는 방문했으니 후보 리스트에서 제거
        if current_node in candidates:
            candidates.remove(current_node)

        for next_node in candidates:
            dfs(next_node, sheep, wolf, candidates.copy())

    dfs(0, 0, 0, [0])
    return max_sheep