[프로그래머스] 양과 늑대 풀이 (2022 카카오 코딩테스트)


문제 설명

2진 트리 모양 초원의 각 노드에 늑대와 양이 한 마리씩 놓여 있습니다. 이 초원의 루트 노드에서 출발하여 각 노드를 돌아다니며 양을 모으려 합니다. 각 노드를 방문할 때 마다 해당 노드에 있던 양과 늑대가 당신을 따라오게 됩니다. 이때, 늑대는 양을 잡아먹을 기회를 노리고 있으며, 당신이 모은 양의 수보다 늑대의 수가 같거나 더 많아지면 바로 모든 양을 잡아먹어 버립니다. 당신은 중간에 양이 늑대에게 잡아먹히지 않도록 하면서 최대한 많은 수의 양을 모아서 다시 루트 노드로 돌아오려 합니다

1

예를 들어, 위 그림의 경우(루트 노드에는 항상 양이 있습니다) 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번 노드는 항상 루트 노드입니다.


풀이

카카오에서 꼭 한 문제씩은 트리 그림이 나오는 것 같다. 하지만 풀이에는 대게 DP가 쓰이고, 본 문제도 마찬가지다.

먼저 문제에서 파악해야할 기본적인 조건들은 방문했던 노드들을 다시 방문할 수 있다는 점양이 한 번이라도 잡아먹히면 안된다는 것이다. 이진 트리가 떡하니 주어진 것 치고는 생소한 조건들이라 접근이 어려울 수 있다. 때문에 예시에서 노드의 방문 순서를 달리해보면서 감을 잡았을 것이다. 하지만 의외로 본 문제의 풀이에서 가장 중요한 핵심은 다음과 같다.

노드를 방문하는 순서는 중요하지 않다!

얼핏보면 노드의 방문 순서가 매우 중요한 것처럼 보인다. 하지만 이 트리의 모든 노드들은 어떻게 해서든 방문할 수 있는 노드절대로 방문할 수 없는 노드 두 가지로 나뉜다. 어떻게 해서든 방문할 수 있는 노드는 노드를 최적의 순서로 이곳 저곳 들려서 양을 최대한 끌어모아서라도 방문할 수 있는 노드이고, 절대로 방문할 수 없는 노드는 아무리 양을 끌어모아도 갈 수 없는 노드를 말한다. 따라서 특정 노드까지 도착할 때 어떤 순서로 왔는지까지는 중요하지 않고 어떤 노드들을 방문해왔는지가 중요하다. 우리는 항상 최적의 수(양이 최대인 경로)만 취급할 것이기 때문이다.

이 아이디어를 통해 노드 별로 현재까지 방문한 노드에 따른 양의 최댓값을 구할 수 있다. 다음 노드를 탐색할 때 양의 최댓값보다 늑대 수가 많아지면 탐색을 종료하고, 그렇지 않다면 계속한다.

set을 선언하여 방문한 노드들을 저장하고 이를 DP에 사용할 수도 있지만, 노드의 최대 갯수가 17개이므로 비트마스킹을 활용하면 좀 더 간단하게 메모이제이션을 할 수 있다. 현재 방문 중인 노드와 방문했던 노드들의 비트를 매개변수로 하는 재귀함수를 통해 top-down 메모이제이션을 진행한다. 이때 방문했던 노드들을 통해 양과 늑대의 개수를 세서 양의 최대값을 갱신함과 동시에 탐색의 지속 여부를 결정한다.

다른 풀이들을 보니 의외로 탐색으로도 해결이 가능한 문제였다. 노드의 최대 갯수가 17개이긴 해도 이미 방문한 노드들을 방문할 수 있다는 점과 실제 코딩테스트에서 5번 문제였다는 점에서 완전탐색으로 해결할 수 있을 것이라는 생각을 못했다. 실제 나올 수 있는 경우의 수가 그리 크지 않은 것 같다. 앞으로는 뒷 번호에서도 충분히 완전탐색 풀이가 가능하다는 것을 인지해야겠다. 여담으로 이번 2022 카카오 코딩테스트 마지막 문제도 결국엔 완전탐색 풀이가 사용된다. 이번에는 카카오가 완전탐색에 꽃혔나보다.


전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
#include <string>
#include <vector>
#include <cstring>
using namespace std;

bool adj[17][17];
int cache[17][(1<<17)-1]; // bitmasking memoization

int solve(int n, int visited, vector<int>& info) {
    int& ret = cache[n][visited];
    if(ret != -1) return ret;

    int sheep = 0, wolf = 0;
    for(int i=0; i<17; i++) {
        if(visited & (1<<i)) {
            if(info[i]) wolf++;
            else sheep++;
        }
    }

    if(sheep == wolf) return ret = 0; // 양이 잡하먹히므로 진행 불가

    ret = sheep;
    for(int i=0; i<17; i++) {
        if(!adj[n][i]) continue;
        ret = max(ret, solve(i, visited | (1<<i), info));
    }

    return ret;
}

int solution(vector<int> info, vector<vector<int>> edges) {
    for(auto v:edges) adj[v[0]][v[1]] = adj[v[1]][v[0]] = true;

    memset(cache, -1, sizeof(cache));
    return solve(0, 1, info);
}
0%