BOJ 12995 트리나라

N개의 노드로 이루어진 트리에서 서로 연결되어 있는 k개의 노드를 고르는 문제이다.
tree dynamic programming 유형이다. BOJ 캠프 문제인데, 아마 tree dp의 기본이 되는 문제일 것이다. 그 당시 백준님의 풀이를 보고 풀었는데, 지금 보니 기억이 안나서 복습하려고 한다. 다시 풀어보고 혼자 생각해 봐야겠다.

d[node][chosen][k] = node를 chosen(0이면 선택 안 함, 1이면 선택)했을 경우에 k개의 노드를 고르는 방법의 수

d[node][0][k] = SUM( d[child][0 or 1][k, k-1] );
d[node][1][k] = SUM( d[child][1][k-1] );
이 될 것 같다. node와 k는 최대 50 이므로 시간 복잡도는 최대  50*50*2*(50*2)..정도로 충분해 보인다... 구현해 봐야겠다. 구현해서 예제를 돌려보니까 알겠다. 이렇게 해서는 안된다. 이런 식으로 하면 위에서 아래로 가면서 선택하는 경우만 세게 된다. 그렇다고 모든 node에서 다 시작해보면 겹치는 것이 생기게 되니까 모든node에서 다 시작하되, 무조건 그 노드를 선택하는 식으로 해도 겹치는 경우가 생길 뿐 아니라... 가장 중요한 것은 한 노드에서 여러 갈래로 갈라지는 것은 커버할 수 없다...

한 노드에서 시작해서, 선택할 경우, 이쪽 자식으로는 몇 개, 저 쪽 자식으로는 몇 개를 선택할 건지에 대해 가능한 경우를 다 해봐야 할 것 같은데 그럼 선택할 수 있는 경우가 너무 많아진다.... 아 그렇기 때문에 백준님의 풀이에서도 이렇게 하지 않는다.

백준님 풀이를 보고 다시 공부해야겠다.
 내가 위에서 시도했던 방법의 경우, 한 노드를 선택한 후 그 아래 자식 노드들을 선택할 수 있는 경우가 너무 많기 때문에 문제를 두 개의 부분으로 나눠서 접근해보자.

일단 입력으로 주어지는 트리는 그래프 형태로 주어지기 때문에, 어느 한 노드를 root로 잡고 트리로 재구성을 해준다.(혹은 parent배열을 만들어서 사용해도 트리로 구성한 효과가 있지만, 트리로 재구성을 하면 자식의 개수도 좀 더 편하게 파악할 수 있고...이 문제에서는 트리로 재구성하는 것이 훨씬 편하다.) 그러면 문제에서 구하는 경우는 어떤 경우이든 간에 트리의 일부가 되기 때문에 어떤 노드를 root로 하는 sub tree에서 그 root를 포함시켜 k개의 노드를 고르는 방법의 수를 각 노드에 대해 구하면 모든 경우를 커버할 수 있고 겹치는 경우도 없을 것이다. 결국 각 노드를 root로 하는 sub tree에서의 경우의 수들을 구해서 다 합하면 답이 된다.

그렇다면 어떤 노드를 root로 하는 sub tree에서의 경우의 수를 어떻게 구할 것인가? 이제 두 개의 부분으로 나눠서 접근해보자. 바로 node의 자식 중 맨 왼쪽 자식을(편의상... 그냥 어느 한 자식이기만 하면 된다) root로 하는 sub tree와 그 나머지로 나눌 수 있다.
그 나머지인 subtree는 node가 root인데, 맨 왼쪽 자식 하나를 root로 하는 sub tree가 분리되어 나갔으므로 원래 tree와는 좀 다르다. 그래서 d[node][e][k]로 표현을 해주는데,
d[node][e][k] = node가 루트일 때, node를 포함해서 k개의 노드를 선택하는 방법의 수 (왼쪽부터 e개의 자식을 root로 하는 sub tree들을 제외한 sub tree에서)
이 된다.

d[node][e][k] = SUM( i = 0~k-1 , d[leftmost child][0][i]*d[node][e+1][k-i] ) 가 될 것이다. 여기서 주의해야 할 것이, node는 무조건 선택해야 하므로 k-i가 0이 되지 않도록 i는 최대 k-1까지이다. 그리고 i가 0일 때는 child를 선택하지 않는 다는 것인데 경우의 수가 곱셈으로 이루어지므로 1을 return해주면 된다.

//코드

//baekjoon님 풀이 슬라이드+ baekjoon님 코드 보고 작성한 코드

#include <cstdio>

#include <cstring>

#include <vector>

using namespace std;

typedef long long ll;

void makeTree(int, int);

ll treeCity(int, int, int);

const ll MOD=1e9+7;

vector<int> adj[51];

vector<int> tree[51];

ll d[51][50][51];

int main() {

    int n, k;

    scanf("%d %d", &n, &k);

    for(int i=0; i<n-1; i++) {

        int u, v;

        scanf("%d %d", &u, &v);

        adj[u].push_back(v);

        adj[v].push_back(u);

    }

    makeTree(1, -1);

    memset(d, -1, sizeof(d));

    ll ans=0;

    for(int i=1; i<=n; i++) {

        ans+=treeCity(i, 0, k);

        ans%=MOD;

    }

    printf("%lld\n", ans);

    return 0;

}

ll treeCity(int root, int e, int k) {

    if(k==0) return 1; //root를 선택하지 않는 경우 return 1

    if(e==tree[root].size()) { //선택할 것이 root만 남은 경우!(자식이 없는 경우도 마찬가지)

        return k==1;

    }

    ll& ret=d[root][e][k];

    if(ret!=-1) return ret;

    ret=0;

    int child=tree[root][e];

    for(int i=0; i<k; i++) {

        //맨 처음의 root는 무조건 선택해야 하므로 k-i가 0이되면 안된다!***

        //맨 처음 root외에 root들은(즉 맨 처음 root의 자식들) 선택되지 않아도 된다.

        //그리고 그 경우(root의 child가 선택되지 않는 경우)는 위쪽에서 k==0일 때, return 1;

        ret+=(treeCity(child, 0, i)*treeCity(root, e+1, k-i));

        ret%=MOD;

    }

    return ret;

}

void makeTree(int root, int parent) {

    for(int i=0; i<adj[root].size(); i++) {

        if(adj[root][i]!=parent) {

            tree[root].push_back(adj[root][i]);

            makeTree(adj[root][i], root);

        }

    }

}

그리고 코드를 다시 짜보면서 실수한 것이 가장 왼쪽의 child를 선택할 때, d[node][0]을 가장 왼쪽 노드로 봤는데, 이러면 처음과 왼쪽 sub tree에서만 맞고, 오른쪽 sub tree에서는 맞지 않다. 그렇다면 어떻게? 바로 e를 이용하는 것이다. 자식 중에서 왼쪽 e개를 제외한 것이므로 맨 왼쪽 자식은 e번째 (0번째 자식부터 시작할 때) 자식이 된다!

그리고 node가 50개 밖에 안되서 k개 뽑는 경우의 수가 int형으로 충분히 커버가능할 줄 알았는데 생각해보니 k값도 최대 50까지 가능하므로 k값에 따라 k개를 뽑는 경우의 수가 매우 커질 수 있다. 그렇기 때문에 long long을 사용해주고, MOD로 꼭 modulo연산을 해줘야 한다.

출처 : 백준님 강의, 강의 자료