[KAKAO BLIND RECRUITMENT / 2022 / Python] 양과 늑대

문제 출처

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번 노드는 항상 루트 노드입니다.

---

입출력 예

---

설계

문제를 처음 접했을 때, 각 노드에서 최대한 많은 양을 모으는 것이 목표임을 고려하여 깊이 우선 탐색(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

    # DFS 탐색 함수
    def dfs(node, sheep, wolf, next_nodes):
        nonlocal max_sheep

        # 현재 노드 방문
        if info[node] == 0:
            sheep += 1
        else:
            wolf += 1

        # 늑대가 양의 수를 초과하면 종료
        if wolf >= sheep:
            return

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

        # 현재 탐색 가능한 노드들 업데이트
        new_next_nodes = next_nodes + graph[node]
        new_next_nodes.remove(node)

        # 다음 노드 탐색
        for next_node in new_next_nodes:
            dfs(next_node, sheep, wolf, new_next_nodes)

    # 초기 상태: 루트 노드(0), 양=0, 늑대=0, 탐색 가능한 노드=[0]
    dfs(0, 0, 0, [0])
    return max_sheep

---

결론 및 느낀점

이 문제를 통해 복잡한 조건을 가진 그래프 탐색 문제를 어떻게 접근하고 해결할 수 있는지에 대해 깊게 생각해 볼 수 있었다. DFS를 활용하여 각 노드를 방문하고, 양과 늑대의 수를 신중하게 관리하면서 탐색을 진행하는 방법을 구현함으로써, 문제에서 요구하는 최적의 결과를 도출할 수 있었다.

특히, 늑대의 수가 양의 수보다 많아지지 않는 경로를 선택하는 과정에서 조건을 정확하게 설정하는 것의 중요성을 다시 한번 깨닫게 되었다. 이번 경험을 통해 더 복잡한 그래프 탐색 문제에도 자신감을 가지고 접근할 수 있는 기반을 마련할 수 있었다.