[백준] 트리나라 (C++)

Yookyubin·2023년 4월 15일
0

백준

목록 보기
13/68

문제

트리나라는 N개의 도시로 이루어져 있고, 각각의 도시는 1번부터 N번까지 번호가 매겨져 있다. 트리나라의 도로 체계는 트리를 이룬다. 즉, 트리나라에는 N-1개의 양방향도로가 있다. 또, 모두 연결되어 있기 때문에, 임의의 두 도시 사이를 항상 오갈 수 있다.

스타트링크의 직원 K명은 트리나라로 이사를 가려고 한다. 모든 직원은 서로 다른 도시로 이사를 가야한다. 즉, 이사할 도시 K개를 선택해야 한다. 이사할 도시에는 중요한 조건이 하나 있는데, 모든 직원이 사는 도시는 연결되어 있어야 한다는 점이다. 예를 들어, 임의의 두 직원 사는 도시가 i와 j라면, i와 j를 연결하는 경로상에 있는 도시에도 직원이 살고 있어야 한다는 점이다.

트리나라의 트리 구조가 주어졌을 때, 이사할 도시 K개를 고르는 방법의 수를 구하는 프로그램을 작성하시오.

입력
첫째 줄에 도시의 수 N과 스타트링크 직원의 수 K가 주어진다. (2 ≤ N ≤ 50, 1 ≤ K ≤ N)
둘째 줄부터 N-1개의 줄에는 도로 정보가 주어진다.

출력
첫째 줄에 도시 K개를 선택하는 방법의 수를 1,000,000,007로 나눈 나머지를 출력한다.

문제 링크

풀이

N개의 노드로 이루어진 트리에서 k개의 이어진 노드를 찾는 문제이다.
즉, k개의 서브트리를 찾으라는 문제가 된다.

처음에는 비트마스킹으로 해결할 수 있나? 생각을 했지만 깊이 생각하기도 전에 노드의 수가 50개나 된다는 것을 보고 포기했다.

그렇다면 어떻게 수많은 DFS 탐색 결과들을 분명하게 저장할 수 있을지 고민해야 했다.

다른 풀이들을 찾아보니 i를 root로하는 트리에서 k개의 노드를 가지는 서브트리의 개수를 dp[i][k]로 표현했다. 하지만 이렇게로는 충분하지 않고 더 분명하고 세분화된 정보를 저장하기 위해 dp 테이블을 한 차원 더 사용하였다.

루트를 포함하여 k개의 노드를 가지는 서브트리의 개수를 구하기 위해 자식서브트리를 임의의 하나와 그 나머지로 나누어서 계산한다.

예를 들어 k=10이고 root -> a, root -> b, root -> c, root -> d로 루트의 자식 서브트리들이 있다면{a}루트 노드를 포함한 {b,c,d} 로 나눈다. 두 개의 집합으로 나누어서 경우의 수를 계산한다.
만약 서브트리 a에서 노드의 개수가 3개인 서브트리를 구한다면 그 나머지 root, {b, c, d}에서 노드의 개수가 7(k-3)인 서브트리를 구한다.
{a}노드의 개수가0부터 k-1까지인 경우를 모두 구하고 root, {b, c, d}에서 k부터 0개의 노드를 가지는 서브트리의 수를 계산한다.
{a}root, {b, c, d}의 경우의 수는 동시에 일어나지만 독립적으로 일어난다. 따라서 두 개의 계산결과를 곱해주면 아무것도 배제하지 않은 첫 경우의 수가 나온다.

그 다음은 루트노드에서 {a}를 제외한 경우를 계산한다.
{b, c, d}에서 다시 임의의 하나와 그 나머지로 나눈다. {b}{c, d}로 나눌 수 있다.
다시 {b}가 0부터 k-1개의 노드를 가지는 서브트리의 개수를 구한다. 나머지도 아까와 같이 root, {c, d}가 k~0개의 노드를 가지는 서브트리의 개수를 구하여 곱한다.
이는 루트에서 a 서브트리를 완전히 배제한 경우의 수이다.

그 다음은 마찬가지로 c를 배제한 경우, d를 배제한 경우를 모두 구해 더해주면
루트노드가 root인 트리의 루트를 포함하여 k개의 노드를 가지는 서브트리의 개수를 구할 수 있다.

위의 과정을 수식으로 표현하기 위해
exclusion번째 서브트리를 {a}라고 한다면, a를 루트로하는 서브트리에서 i개의 노드를 가지는 서브트리의 개수는

dp[a][0][i]dp[a][0][i]

나머지는

dp[root][exclusion+1][ki]dp[root][exclusion+1][k-i]

따라서 루트 노드가 root이고 exclusion 이전의 서브트리들은 배제한 상태의 k개의 노드를 갖는 서브트리의 경우의 수는

dp[root][exclusion][k]=Σdp[a][0][i]dp[root][exculsion+1][ki]  (i:  0    k1)dp[root][exclusion][k] = \Sigma{dp[a][0][i] * dp[root][exculsion+1][k-i]} \;_{(i: \; 0\; \sim\;k-1)}

이라는 식이 완성된다.

나머지 부분을 계산할때 k-i가 0이 되어버리면 루트노드를 포함하지 않는 경우가 되어버려 k-i는 0이되면 안된다. 따라서 반복문을 돌때에도 i=0 부터 k-1까지 반복한다.

코드

#include <iostream>
#include <vector>

using namespace std;

int n, k;
vector<vector<int>> graph;
vector<vector<int>> tree;
vector<vector<vector<int>>> dp;
const int64_t MOD = 1e9 + 7;

void makeTree(int current, int parent){
    for (int child : graph[current]){
        if (child == parent) continue;

        tree[current].push_back(child);
        makeTree(child, current);
    }
}

int64_t countSubtree(int root, int exclusion, int kount){
    if (kount == 0) return 1;
    if (exclusion == tree[root].size()) return kount==1; // 루트만 남았는데 세야할 노드 수가 1이 아닌경우
    if (dp[root][exclusion][kount] != -1) return dp[root][exclusion][kount]; // 이미 계산한 적이 있는 값

    int64_t ret = 0; 
    int child = tree[root][exclusion];
    for (int i=0; i < kount; i++){
        ret += countSubtree(child, 0, i) * countSubtree(root, exclusion+1, kount-i);
        ret %= MOD;
    }
    return dp[root][exclusion][kount] = ret;
}

int main() {
    cin.tie(0);
    ios_base::sync_with_stdio(false);

    cin >> n >> k;
    
    graph.assign(n+1, vector<int>());
    tree.assign(n+1, vector<int>());
    dp.assign(n+1, vector<vector<int>>(n+1, vector<int>(n+1, -1)));

    for (int i=0; i < n-1; i++){
        int a, b;
        cin >> a >> b;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }

    makeTree(1, -1);
    
    int answer = 0;
    for(int i=1; i < n+1; i++){
        answer += countSubtree(i, 0, k);
        answer %= MOD;
    }
    cout << answer << "\n";
    return 0;
}

참고

https://foxtrotin.tistory.com/338
https://fullalgorithmpanic.blogspot.com/2016/10/boj-12995.html
https://viyoung.tistory.com/127

profile
붉은다리 제프

0개의 댓글