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

문제 요약

이진 트리 형태의 초원에 양과 늑대가 있다.

루트 노드에서 시작해 노드를 방문하면서 양을 최대한 많이 모아야 한다.

각 노드를 방문하면 해당 노드의 양 또는 늑대가 따라온다.

단, 어느 순간이라도 늑대 수가 양 수 이상이 되면 양이 모두 잡아먹히므로 그 경로는 더 이상 진행할 수 없다.

목표는 조건을 지키면서 모을 수 있는 양의 최대 수를 구하는 것이다.

핵심 아이디어

이 문제에서 중요한 점은 단순히 트리의 자식 방향으로만 DFS를 하는 것이 아니라는 점이다.

현재까지 방문한 노드들로 인해 갈 수 있게 된 후보 노드들이 여러 개 생길 수 있다.

다음 이동은 현재 노드의 자식 중 하나가 아니라, 지금까지 열린 후보 노드들 중 하나를 선택하는 방식이다.

따라서 DFS 상태에 다음 정보를 들고 다녀야 한다.

현재 양의 수
현재 늑대의 수
다음에 방문 가능한 노드 후보 목록

후보 노드 중 하나를 선택해 방문하고, 그 노드의 자식들을 다시 후보 목록에 추가한다.

트리 구성

edges는 부모와 자식 관계를 담고 있다.

각 노드의 자식 노드를 빠르게 찾기 위해 인접 리스트를 만든다.

children = [[] for _ in range(len(info))]

for parent, child in edges:
    children[parent].append(child)

DFS 상태

DFS 함수는 다음 값을 인자로 받는다.

dfs(sheep, wolf, candidates)

• sheep: 지금까지 모은 양의 수
• wolf: 지금까지 따라온 늑대의 수
• candidates: 다음에 방문할 수 있는 노드 목록

루트 노드는 항상 양이므로 시작 상태는 다음과 같다.

dfs(1, 0, children[0])

후보 노드 선택하기

현재 후보 목록에서 하나의 노드를 선택한다.

선택한 노드를 방문하면 다음 후보 목록은 다음과 같이 바뀐다.

1. 기존 후보 목록에서 선택한 노드를 제거한다.
2. 선택한 노드의 자식들을 후보 목록에 추가한다.

next_candidates = candidates[:]
next_candidates.remove(node)
next_candidates.extend(children[node])

이렇게 하면 지금까지 갈 수 있게 된 노드들을 계속 유지하면서 탐색할 수 있다.

늑대 조건 확인

방문한 노드가 양인지 늑대인지에 따라 개수를 갱신한다.

if info[node] == 0:
    next_sheep += 1
else:
    next_wolf += 1

그 후 늑대 수가 양 수 이상이면 더 진행하지 않는다.

if next_wolf >= next_sheep:
    continue

Python 코드

def solution(info, edges):
    n = len(info)
    children = [[] for _ in range(n)]

    for parent, child in edges:
        children[parent].append(child)

    answer = 0

    def dfs(sheep, wolf, candidates):
        nonlocal answer

        answer = max(answer, sheep)

        for node in candidates:
            next_sheep = sheep
            next_wolf = wolf

            if info[node] == 0:
                next_sheep += 1
            else:
                next_wolf += 1

            if next_wolf >= next_sheep:
                continue

            next_candidates = candidates[:]
            next_candidates.remove(node)
            next_candidates.extend(children[node])

            dfs(next_sheep, next_wolf, next_candidates)

    dfs(1, 0, children[0])

    return answer

코드 설명

자식 노드 목록 만들기

children = [[] for _ in range(n)]

for parent, child in edges:
    children[parent].append(child)

각 노드에서 바로 이동 가능해지는 자식 노드들을 저장한다.

정답 갱신

answer = max(answer, sheep)

DFS의 각 상태에서 현재까지 모은 양의 수로 정답을 갱신한다.

후보 노드 순회

for node in candidates:

현재 갈 수 있는 후보 노드 중 하나를 선택한다.

이 문제는 현재 위치가 중요한 것이 아니라, 지금까지 방문해서 열어둔 후보 노드 집합이 중요하다.

양과 늑대 수 갱신

if info[node] == 0:
    next_sheep += 1
else:
    next_wolf += 1

info[node]가 0이면 양, 1이면 늑대다.

실패 조건

if next_wolf >= next_sheep:
    continue

늑대 수가 양 수 이상이 되면 그 즉시 실패하므로 해당 선택은 더 진행하지 않는다.

후보 목록 갱신

next_candidates = candidates[:]
next_candidates.remove(node)
next_candidates.extend(children[node])

방문한 노드는 후보 목록에서 제거한다.

그리고 방문한 노드의 자식들이 새롭게 이동 가능한 노드가 되므로 후보 목록에 추가한다.

왜 현재 위치를 저장하지 않을까?

문제에서는 루트에서 출발해 돌아다닌다고 표현하지만, 실제 탐색에서는 현재 위치보다 방문 가능한 후보 노드 집합이 더 중요하다.

이미 방문한 노드를 통해 이동할 수 있으므로, 지금까지 열린 노드 중 어떤 노드를 다음에 방문할지 선택할 수 있다.

따라서 DFS 상태에는 현재 위치 하나가 아니라 candidates를 저장한다.

시간 복잡도

노드 수를 N이라고 하면, 가능한 후보 선택 순서를 탐색한다.

최악의 경우 많은 순열 상태를 탐색할 수 있다.

O(탐색 가능한 상태 수)

프로그래머스 문제 제한에서는 노드 수가 작기 때문에 백트래킹으로 충분히 해결할 수 있다.

공간 복잡도

DFS 재귀 깊이는 최대 노드 수만큼 깊어질 수 있다.

후보 목록도 상태마다 복사된다.

O(N^2)

정리

이 문제는 단순 트리 DFS가 아니라, 현재까지 방문해서 갈 수 있게 된 후보 노드들 중 하나를 선택하는 백트래킹 문제다.

핵심은 다음과 같다.

• 루트에서 시작한다.
• 현재 갈 수 있는 후보 노드 목록을 관리한다.
• 후보 중 하나를 방문하면 그 노드의 자식들을 후보에 추가한다.
• 늑대 수가 양 수 이상이면 탐색을 중단한다.
• 탐색 중 모은 양의 최댓값을 정답으로 갱신한다.

후보 목록을 상태로 들고 다니는 것이 이 문제의 가장 중요한 포인트다.