백준 22199번 Cat in a tree

https://codeforces.com/blog/entry/134665

코드포스 블로그를 둘러보다 트리DP에서 노드의 정보를 스몰 투 라지로 합치는 검은돌 트릭이 서브트리의 크기가 아닌 최대 깊이에 비례하는경우 상당히 빨라진다는 글을 보았다. qwerasdfzxcl 님이 말하기로는 시간복잡도가 무려 O(N)이 나온다고 하였다. 최근 팀연습에 나온 요코하마 2023 J번이나(이건 정해가 HLD이다! 검은돌 트릭을 쓰면 1초면 충분하다!) 이번 카이스트 Mock L번 등 최근 자주 사용되는 유형인것 같아 여러 경우에 따른 시간복잡도를 알아두면 꽤나 도움이 될것같다.

 

https://www.acmicpc.net/problem/22199

소개된 연습문제 하나가 백준에 있어 풀어보았다. 내가 짠 코드는 deque가 생각보다 좀 심하게 느린지 log가 붙은듯 했지만 아무튼 O(N^2)보다는 빠르니 별 문제는 없을것같다.

 

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;

vector<vector<int>> graph;
vector<deque<int>> dp;
int D;

void dfs(int node){
    for (int i : graph[node]){
        dfs(i);
        if (dp[node].size() < dp[i].size()) swap(dp[node], dp[i]);
        for (int j = 0; j < dp[i].size(); j++){
            int tdn = dp[node][j];
            dp[node][j] = max(dp[node][j], dp[i][j]);
            if ((int)dp[node].size() > D - j - 2){
                if ((j + 1) << 1 < D) dp[node][j] = max(dp[node][j], dp[node][D - j - 2] + dp[i][j]);
                if ((int)dp[i].size() > D - j - 2) dp[node][j] = max(dp[node][j], tdn + dp[i][max(D - j - 2, j)]);
            }
        }
        for (int j = (int)dp[i].size() - 2; j >= 0; j--){
            dp[node][j] = max(dp[node][j], dp[node][j + 1]);
        }
    }
    
    if (dp[node].empty()) dp[node].push_front(1);
    else dp[node].push_front(dp[node].front());
    if (dp[node].size() > D){
        dp[node].front() = max(dp[node].front(), dp[node].back() + 1);
        dp[node].pop_back();
    }
    
    /*cout << node << " : ";
    for (int i : dp[node]){
        cout << i << ' ';
    }
    cout << '\n';*/
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    int N;
    cin >> N >> D;
    
    if (D == 1){
        cout << N << '\n';
        return 0;
    }
    
    graph.resize(N);
    dp.resize(N);
    
    for (int i = 1; i < N; i++){
        int x;
        cin >> x;
        graph[x].push_back(i);
    }
    
    dfs(0);
    cout << dp[0].front() << '\n';
}

 

우선 DP식은 DP[i][j] -> i 번 노드의 서브트리에서 깊이가 j 이하인 노드를 사용하지 않고 칠할수 있는 노드의 개수로 세웠다. 노드를 합칠때는 일단 각 노드의 루트 깊이가 다르면 안되니 현재 노드는 비워둔 상태로 합치고 마지막에 현재 노드를 칠한(다르게 말하면 현재 서브트리의 루트가 칠해진) 경우를 처리하고 D를 넘어가는 필요 없는 경우는 제거하면 된다. 노드를 합칠때는 DP[i1][j] = max(DP[i1][j], DP[i2][j], DP[i1][j] + DP[i2][max(j, D - j - 2)], DP[i1][max(j, D - j - 2)] + DP[i2][j])가 되게 된다. 두 노드의 서브트리간에 조건이 성립하면서 칠해진 가장 깊이가 낮은 노드의 정보를 담고 있어야 하기 때문이다. 이후 배열이 단조 감소하게 유지하며 노드의 DP값을 합칠 수 없는 D = 1인 경우를 예외처리해주면 AC를 받을 수 있다. 풀고나서 기여창을 열어보니 검은돌 트릭 이외에도 센트로이드 분할과 그리디 풀이로도 풀리는 재미있는 문제였다.