BOJ 2454 트리 분할

트리 분할
(http://boj.kr/2454)

KOI 고등부 3번 문제인데 내 힘으로 풀게 되어서 꽤 뿌듯하다.

문제 이름대로 주어지는 그래프는 트리형태이니 루트를 기준으로 내려간다고 생각해보자.
그 뒤 루트가 \(i\)번 노드인 서브트리에서의 답을 \(A_i\)라고 하자.
그리고 그 답에서의 \(i\)번 노드를 포함하는 경로의 최소 크기를 \(B_i\)라고 하자.
그러면 이 두 가지를 가지고 DP로 답을 찾을 수 있다.

현재 노드가 \(u\)일 때, \(u\)가 루트인 서브트리에서 답이 될 수 있는 경우는 3가지가 있다.

1. \(u\)의 자식들의 답은 그대로 두고, 그 상태에서 \(u\) 하나만 있는 경로를 추가하는 경우
2. \(u\)의 자식 중 하나의 경로에 \(u\)를 추가하는 경우
3. \(u\)의 자식들 중 두 개를 \(u\)를 포함시켜 잇는 경우

이 때, 1은 문제가 없지만 2, 3은 당연히 그 경우에 경로의 길이가 \(K\)를 넘지 않아야 한다.

각 경우에 \(A_u\)와 \(B_u\)는,

1. \(A_u\)는 \(\text{자식들의 }A\text{의 합}+1\)이고 \(B_u\)는 \(0\)이다.
2. 만약 자식 중에 \(K\)보다 작은 \(B\)가 존재한다면, \(A_u\)는 \(\text{자식들의 }A\text{의 합}\)이고 \(B_u\)는 \(\text{자식들의 }B\text{의 최소값}+1\)이다.
3. 만약 서로 다른 두 자식의 \(B\text{의 합}\)이 \(K\)보다 작다면, \(A_u\)는 \(\text{자식들의 }A\text{의 합}-1\)이고 \(B_u\)는 \(K\)로 처리한다. (왜냐하면, 이렇게 이었을 때 \(u\)의 부모는 이 경로에 포함될 수 없다.)

이 3가지 경우를 잘 처리해주면 된다.

#include <bits/stdc++.h>
 
using namespace std;
 
int k;
vector<int> gph[300001];
pair<int, int> d[300001];
 
void f(int now, int prv) {
    vector<int> v;
    int tot = 0;
    for (int nxt : gph[now])
        if (nxt != prv) {
            f(nxt, now);
            tot += d[nxt].first;
            v.push_back(d[nxt].second);
        }
 
    d[now] = { tot + 1, 1 };
 
    for (int nxt : gph[now])
        if (nxt != prv && d[nxt].second <= k)
            d[now] = min(d[now], { tot, d[nxt].second + 1 });
 
    sort(v.begin(), v.end());
    if (v.size() >= 2 && v[0] + v[1] <= k)
        d[now] = min(d[now], { tot - 1, k + 1 });
}
 
int main() {
    int n;
    scanf("%d %d", &n, &k);
    for (int i = 0;i < n - 1;++i) {
        int u, v;
        scanf("%d %d", &u, &v);
        gph[u].push_back(v);
        gph[v].push_back(u);
    }
 
    f(1, 0);
 
    printf("%d", d[1].first);
 
    return 0;
}