센트로이드 분할

트리를 무게중심으로 쪼개 경로 문제를 분할 정복.

트리

플래티넘 I

선수 지식: 트리 DP

1강 무게중심으로 분할 정복 공식

트리 위의 분할 정복

배열의 분할 정복은 "절반으로 쪼개 각각 풀고, 경계를 가로지르는 경우만 따로
합친다." 트리에서 이를 흉내 내는 것이 센트로이드 분할(centroid
decomposition) 이다. "경로 길이가 정확히 $$K$$ 인 쌍의 수" 같은 모든 정점
쌍/경로 문제에 강력하다.

무게중심(센트로이드)

트리의 무게중심은 그 정점을 제거했을 때 남는 모든 조각의 크기가
$\lfloor N/2 \rfloor$ 이하가 되는 정점이다. 모든 트리에는 무게중심이 (1개
또는 2개) 항상 존재한다.

핵심 성질: 무게중심을 기준으로 트리를 쪼개면, 각 조각의 크기가 절반 이하다.
그래서 재귀 깊이가 $O(\log N)$ 이다.

무게중심을 $$c$$ 라 하면, 트리의 모든 경로는 (1) $$c$$ 를 지나거나 (2) $$c$$ 를 떼어
낸 어떤 한 조각 안에 완전히 들어간다. (1)을 $$c$$ 에서 한 번에 처리하고, (2)는 각
조각에 대해 재귀하면 모든 경로를 빠짐없이 한 번씩 본다.

무게중심 찾기

한 번의 DFS로 각 정점의 서브트리 크기 sz를 구한다.
현재 조각 크기를 total이라 할 때, 어떤 정점에서 (가장 큰 자식 조각 크기)와
( $total - sz[\text{현재}]$ , 즉 부모 쪽 조각)가 모두 $\le total/2$ 이면 그
정점이 무게중심이다.

$O(\text{조각 크기})$ 에 찾는다.

분할 정복 절차

현재 조각의 무게중심 $$c$$ 를 찾는다.
$$c$$ 를 "삭제됨"으로 표시하고, $$c$$ 를 지나는 경로의 기여를 계산한다(보통
각 조각에서 $$c$$ 까지의 거리를 모아 카운팅).
$$c$$ 를 떼어 낸 각 인접 조각에 대해 재귀한다.

복잡도

재귀 깊이가 $O(\log N)$ 이고, 각 깊이에서 모든 정점이 정확히 한 번씩 어떤
무게중심의 조각에 속한다. 깊이마다 처리 비용이 $$O(N)$$ (또는 정렬·이분 탐색이면
$O(N \log N)$ )이라 전체 $O(N \log N)$ 또는 $O(N \log^2 N)$ 이다.

핵심 보장: 한 정점은 $O(\log N)$ 개의 무게중심 조각에만 속한다. 이 사실이
모든 센트로이드 분할 풀이의 복잡도 근거다.

무엇을 푸나

"거리가 $\le K$ 인 정점 쌍 수", "거리가 정확히 $$K$$ 인 쌍 수", 그리고 동적으로
정점에 색을 칠하며 "가장 가까운 특정 색 정점까지 거리"를 묻는 문제(센트로이드
트리에 보조 정보 저장)까지 다룬다.

2강 구현과 활용 공식

거리가 K 이하인 쌍의 수

각 무게중심에서, 서로 다른 조각에 속한 두 정점까지의 거리 합이 $\le K$ 인 쌍을
센다. 같은 조각 안 쌍은 그 조각의 무게중심을 거치지 않으므로 포함-배제로
빼 준다.

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 100001;
vector<int> g[MAXN];
bool removed[MAXN];
int sz[MAXN];
int n, K;
long long answer = 0;

int calc_size(int u, int p) {       // 조각 크기
    sz[u] = 1;
    for (int v : g[u]) if (v != p && !removed[v]) sz[u] += calc_size(v, u);
    return sz[u];
}
int find_centroid(int u, int p, int total) {
    for (int v : g[u]) if (v != p && !removed[v])
        if (sz[v] > total / 2) return find_centroid(v, u, total);
    return u;
}
void collect(int u, int p, int d, vector<int>& ds) {  // 거리 모으기
    ds.push_back(d);
    for (int v : g[u]) if (v != p && !removed[v]) collect(v, u, d + 1, ds);
}
long long count_pairs(vector<int>& ds) {   // 합 <= K 인 쌍 수
    sort(ds.begin(), ds.end());
    long long res = 0;
    int lo = 0, hi = ds.size() - 1;
    while (lo < hi) {                       // 투 포인터
        if (ds[lo] + ds[hi] <= K) { res += hi - lo; lo++; }
        else hi--;
    }
    return res;
}
void decompose(int entry) {
    calc_size(entry, -1);
    int c = find_centroid(entry, -1, sz[entry]);
    removed[c] = true;
    vector<int> all = {0};                  // 무게중심 자신(거리 0)
    for (int v : g[c]) if (!removed[v]) {
        vector<int> ds;
        collect(v, c, 1, ds);
        answer -= count_pairs(ds);          // 같은 조각 쌍은 포함-배제로 제거
        for (int d : ds) all.push_back(d);
    }
    answer += count_pairs(all);             // c를 지나는 모든 쌍
    for (int v : g[c]) if (!removed[v]) decompose(v);
}

decompose(0) 후 answer가 거리 $\le K$ 인 정점 쌍의 수다.

흔한 실수

무게중심 표시 후 크기 재계산 — 한 조각으로 재귀할 때마다 calc_size를
다시 호출해야 한다. 부모 쪽으로 새지 않도록 removed를 항상 확인한다.
포함-배제 누락 — 같은 조각 안 쌍(무게중심 거리 합이 잘못 세진 것)을 빼지
않으면 중복·오답이 난다.
무게중심 자신 포함 — all에 거리 0(무게중심 자신)을 넣어야 " $$c$$ 와 다른
정점" 쌍도 센다.
재귀 깊이 — 분할 깊이는 $O(\log N)$ 이지만 collect/calc_size의 DFS는
깊어질 수 있으니 큰 입력은 스택에 주의한다.

활용

문제	무게중심에서 모으는 정보
거리 $\le K$ 쌍 수	거리 배열 + 투 포인터
거리 정확히 $$K$$ 쌍 수	거리별 빈도 배열
동적 "가장 가까운 색칠 정점"	센트로이드 트리에 각 조상까지 거리 저장

센트로이드 분할은 "트리의 모든 경로를 $O(\log N)$ 겹으로 분해"한다는 한 가지
아이디어로 매우 다양한 경로·쌍 문제를 통일한다. 센트로이드 트리를 명시적으로
세워 보조 자료 구조를 얹는 동적 버전까지 익히면 응용 폭이 크게 넓어진다.

문제	무게중심에서 모으는 정보
거리 $\le K$ 쌍 수	거리 배열 + 투 포인터
거리 정확히 $\(K\)$ 쌍 수	거리별 빈도 배열
동적 "가장 가까운 색칠 정점"	센트로이드 트리에 각 조상까지 거리 저장