트리 체인 분할 개요
트리 체인 분할(Heavy Path Decomposition)은 트리 구조에서 효율적인 쿼리 처리를 위한 고급 자료구조 기법이다. 이 기법은 다음 네 가지 핵심 연산을 지원한다:
- 두 노드 x에서 y까지의 최단 경로상의 모든 노드에 값을 더한다
- 두 노드 x에서 y까지의 최단 경로상의 모든 노드 값의 합을 구한다
- 노드 x를 루트로 하는 서브트리의 모든 노드에 값을 더한다
- 노드 x를 루트로 하는 서브트리의 모든 노드 값의 합을 구한다
핵심 아이디어는 복잡한 트리 구조를 여러 개의.disjoint한 체인으로 분할한 후, 각 체인을 구간 데이터 струкures로 관리하는 것이다.
기본 정의
헤비 차일드(Heavy Child): 특정 노드의 모든 자식 중 서브트리 크기가 가장 큰 자식을 의미한다.
헤비 엣지(Heavy Edge): 노드와 그 헤비 차일드를 연결하는 엣지이다.
헤비 체인(Heavy Path): 헤비 엣지들을 연속적으로 연결하여形成되는 체인이다.
라이트 체인(Light Path): 헤비 엣지가 아닌 일반 엣지들로 연결된 체인이다.
이러한 정의를 활용하면, 각 노드의 헤비 차일드를 우선적으로 연결함으로써 자연스럽게 트리를 여러 개의 헤비 체인과 라이트 체인으로 분할할 수 있다.
트리를 체인으로 분할하기
분할된 각 체인 내의 노드 번호가 연속적이어야 구간 쿼리가 가능하다. 따라서 DFS 순서에 따라 트리 전체에 새로운 번호를 매기는 작업이 필수적이다.
중요한 점은 헤비 체인 내의 노드 번호가 연속적으로 유지되도록 헤비 차일드를 우선적으로 방문해야 한다는 것이다.
이를 위해 두 번의 DFS를 수행한다:
- 첫 번째 DFS: 각 노드의 서브트리 크기, 깊이, 부모, 헤비 차일드를 계산
- 두 번째 DFS: 헤비 체인 내 노드들이 연속 번호를 갖도록 새로운 번호 할당
코드 구현
변수 선언
int depth[MAXN]; // 노드의 깊이
int parent[MAXN]; // 노드의 부모
int heavy[MAXN]; // 노드의 헤비 차일드
int subtreeSize[MAXN]; // 노드의 서브트리 크기
int chainHead[MAXN]; // 노드가 속한 체인의 헤드
int dfsOrder[MAXN]; // DFS 순서상 해당 노드의 번호
int nodeValue[MAXN]; // 각 노드의 원래 가중치
첫 번째 DFS: 기본 정보 계산
int dfsFirst(int current, int parentNode, int depthValue) {
depth[current] = depthValue;
parent[current] = parentNode;
subtreeSize[current] = 1;
int maxSize = -1;
for (Edge* e = adj[current]; e != nullptr; e = e->next) {
if (e->to == parentNode) continue;
subtreeSize[current] += dfsFirst(e->to, current, depthValue + 1);
if (subtreeSize[e->to] > maxSize) {
maxSize = subtreeSize[e->to];
heavy[current] = e->to;
}
}
return subtreeSize[current];
}
두 번째 DFS: 번호 재할당
void dfsSecond(int current, int chainTop) {
dfsOrder[current] = ++timestamp;
segArray[timestamp] = nodeValue[current];
chainHead[current] = chainTop;
if (heavy[current] != 0) {
dfsSecond(heavy[current], chainTop);
}
for (Edge* e = adj[current]; e != nullptr; e = e->next) {
if (dfsOrder[e->to] == 0) {
dfsSecond(e->to, e->to);
}
}
}
세그먼트 트리 구축
struct SegmentNode {
int left, right;
int size;
int value;
int lazy;
};
SegmentNode seg[MAXN];
void buildSegment(int idx, int l, int r) {
seg[idx].left = l;
seg[idx].right = r;
seg[idx].size = r - l + 1;
if (l == r) {
seg[idx].value = segArray[l];
return;
}
int mid = (l + r) >> 1;
buildSegment(idx << 1, l, mid);
buildSegment(idx << 1 | 1, mid + 1, r);
seg[idx].value = (seg[idx << 1].value + seg[idx << 1 | 1].value) % MOD;
}
세그먼트 트리 기본 연산
void applyLazy(int idx, int addValue) {
seg[idx].value = (seg[idx].value + seg[idx].size * addValue) % MOD;
seg[idx].lazy = (seg[idx].lazy + addValue) % MOD;
}
void pushDown(int idx) {
if (seg[idx].lazy == 0) return;
applyLazy(idx << 1, seg[idx].lazy);
applyLazy(idx << 1 | 1, seg[idx].lazy);
seg[idx].lazy = 0;
}
void rangeAdd(int idx, int l, int r, int value) {
if (l <= seg[idx].left && seg[idx].right <= r) {
applyLazy(idx, value);
return;
}
pushDown(idx);
int mid = (seg[idx].left + seg[idx].right) >> 1;
if (l <= mid) rangeAdd(idx << 1, l, r, value);
if (r > mid) rangeAdd(idx << 1 | 1, l, r, value);
seg[idx].value = (seg[idx << 1].value + seg[idx << 1 | 1].value) % MOD;
}
int rangeQuery(int idx, int l, int r) {
if (l <= seg[idx].left && seg[idx].right <= r) {
return seg[idx].value;
}
pushDown(idx);
int mid = (seg[idx].left + seg[idx].right) >> 1;
int result = 0;
if (l <= mid) result = (result + rangeQuery(idx << 1, l, r)) % MOD;
if (r > mid) result = (result + rangeQuery(idx << 1 | 1, l, r)) % MOD;
return result;
}
트리의 경로 연산 구현
핵심 원리: 서로 다른 헤비 체인에 있는 두 노드는各自的 체인으로 올라가면서 같은 체인에 도달할 때까지 반복한다.
void pathQuery(int u, int v) {
int answer = 0;
while (chainHead[u] != chainHead[v]) {
if (depth[chainHead[u]] < depth[chainHead[v]]) swap(u, v);
answer = (answer + rangeQuery(1, dfsOrder[chainHead[u]], dfsOrder[u])) % MOD;
u = parent[chainHead[u]];
}
if (depth[u] > depth[v]) swap(u, v);
answer = (answer + rangeQuery(1, dfsOrder[u], dfsOrder[v])) % MOD;
printf("%d\n", answer);
}
void pathUpdate(int u, int v, int value) {
while (chainHead[u] != chainHead[v]) {
if (depth[chainHead[u]] < depth[chainHead[v]]) swap(u, v);
rangeAdd(1, dfsOrder[chainHead[u]], dfsOrder[u], value);
u = parent[chainHead[u]];
}
if (depth[u] > depth[v]) swap(u, v);
rangeAdd(1, dfsOrder[u], dfsOrder[v], value);
}
서브트리 연산
서브트리 노드들은 dfsOrder 기준으로 연속적으로 위치하므로 간단하게 처리한다.
void subtreeUpdate(int node, int value) {
rangeAdd(1, dfsOrder[node], dfsOrder[node] + subtreeSize[node] - 1, value);
}
int subtreeQuery(int node) {
return rangeQuery(1, dfsOrder[node], dfsOrder[node] + subtreeSize[node] - 1);
}
시간 복잡도 분석
성질 1: 라이트 엣지 (u, v)에서 v의 서브트리 크기는 u의 서브트리 크기의 1/2 이하이다.
성질 2: 두 노드 사이의 경로에서 라이트 엣지의 개수는 O(log n) 이하이다.
経路を查询/更新할 때마다 최대 O(log n)개의 헤비 체인을 거치게 되고, 각 헤비 체인에 대한 구간 연산은 O(log n)이므로 총 시간 복잡도는 O(log² n)이다.
전체 구현
#include <iostream>
#include <cstdio>
#include <cstring>
using namespace std;
const int MAXN = 2000010;
const int MOD = 1000000007;
struct Edge {
int to;
Edge* next;
};
struct SegmentNode {
int l, r, sz;
int val, lazy;
};
int vertexCount, queryCount, rootVertex;
int adj[MAXN];
Edge edges[MAXN * 2];
int edgeCnt = 1;
int depthArr[MAXN];
int parentArr[MAXN];
int heavyChild[MAXN];
int subtreeSize[MAXN];
int chainTop[MAXN];
int dfsIndex[MAXN];
int nodeWeight[MAXN];
int segmentArray[MAXN];
int tmpSize = 0;
SegmentNode segTree[MAXN];
inline void addEdge(int from, int to) {
edges[edgeCnt].to = to;
edges[edgeCnt].next = (Edge*)&adj[from];
*(int*)&adj[from] = edgeCnt++;
}
int dfs1(int v, int p, int d) {
depthArr[v] = d;
parentArr[v] = p;
subtreeSize[v] = 1;
int maxSub = -1;
for (int i = adj[v]; i != -1; i = edges[i].next) {
if (edges[i].to == p) continue;
subtreeSize[v] += dfs1(edges[i].to, v, d + 1);
if (subtreeSize[edges[i].to] > maxSub) {
maxSub = subtreeSize[edges[i].to];
heavyChild[v] = edges[i].to;
}
}
return subtreeSize[v];
}
void dfs2(int v, int top) {
dfsIndex[v] = ++tmpSize;
segmentArray[tmpSize] = nodeWeight[v];
chainTop[v] = top;
if (heavyChild[v] != 0) {
dfs2(heavyChild[v], top);
}
for (int i = adj[v]; i != -1; i = edges[i].next) {
if (dfsIndex[edges[i].to] == 0) {
dfs2(edges[i].to, edges[i].to);
}
}
}
void segUpdate(int idx) {
segTree[idx].val = (segTree[idx*2].val + segTree[idx*2+1].val) % MOD;
}
void segBuild(int idx, int l, int r) {
segTree[idx].l = l;
segTree[idx].r = r;
segTree[idx].sz = r - l + 1;
if (l == r) {
segTree[idx].val = segmentArray[l];
return;
}
int mid = (l + r) >> 1;
segBuild(idx*2, l, mid);
segBuild(idx*2+1, mid+1, r);
segUpdate(idx);
}
void segPush(int idx) {
if (segTree[idx].lazy == 0) return;
segTree[idx*2].val = (segTree[idx*2].val + segTree[idx*2].sz * segTree[idx].lazy) % MOD;
segTree[idx*2+1].val = (segTree[idx*2+1].val + segTree[idx*2+1].sz * segTree[idx].lazy) % MOD;
segTree[idx*2].lazy = (segTree[idx*2].lazy + segTree[idx].lazy) % MOD;
segTree[idx*2+1].lazy = (segTree[idx*2+1].lazy + segTree[idx].lazy) % MOD;
segTree[idx].lazy = 0;
}
void segAdd(int idx, int l, int r, int v) {
if (l <= segTree[idx].l && segTree[idx].r <= r) {
segTree[idx].val = (segTree[idx].val + segTree[idx].sz * v) % MOD;
segTree[idx].lazy = (segTree[idx].lazy + v) % MOD;
return;
}
segPush(idx);
int mid = (segTree[idx].l + segTree[idx].r) >> 1;
if (l <= mid) segAdd(idx*2, l, r, v);
if (r > mid) segAdd(idx*2+1, l, r, v);
segUpdate(idx);
}
int segQuery(int idx, int l, int r) {
if (l <= segTree[idx].l && segTree[idx].r <= r) {
return segTree[idx].val;
}
segPush(idx);
int mid = (segTree[idx].l + segTree[idx].r) >> 1;
int ans = 0;
if (l <= mid) ans = (ans + segQuery(idx*2, l, r)) % MOD;
if (r > mid) ans = (ans + segQuery(idx*2+1, l, r)) % MOD;
return ans;
}
void updatePath(int u, int v, int val) {
while (chainTop[u] != chainTop[v]) {
if (depthArr[chainTop[u]] < depthArr[chainTop[v]]) swap(u, v);
segAdd(1, dfsIndex[chainTop[u]], dfsIndex[u], val);
u = parentArr[chainTop[u]];
}
if (depthArr[u] > depthArr[v]) swap(u, v);
segAdd(1, dfsIndex[u], dfsIndex[v], val);
}
int queryPath(int u, int v) {
int ans = 0;
while (chainTop[u] != chainTop[v]) {
if (depthArr[chainTop[u]] < depthArr[chainTop[v]]) swap(u, v);
ans = (ans + segQuery(1, dfsIndex[chainTop[u]], dfsIndex[u])) % MOD;
u = parentArr[chainTop[u]];
}
if (depthArr[u] > depthArr[v]) swap(u, v);
ans = (ans + segQuery(1, dfsIndex[u], dfsIndex[v])) % MOD;
return ans;
}
int main() {
memset(adj, -1, sizeof(adj));
scanf("%d %d %d %d", &vertexCount, &queryCount, &rootVertex, &MOD);
for (int i = 1; i <= vertexCount; i++) {
scanf("%d", &nodeWeight[i]);
}
for (int i = 1; i < vertexCount; i++) {
int x, y;
scanf("%d %d", &x, &y);
addEdge(x, y);
addEdge(y, x);
}
dfs1(rootVertex, 0, 1);
dfs2(rootVertex, rootVertex);
segBuild(1, 1, vertexCount);
while (queryCount--) {
int op, x, y, z;
scanf("%d", &op);
if (op == 1) {
scanf("%d %d %d", &x, &y, &z);
z %= MOD;
updatePath(x, y, z);
}
else if (op == 2) {
scanf("%d %d", &x, &y);
printf("%d\n", queryPath(x, y));
}
else if (op == 3) {
scanf("%d %d", &x, &z);
z %= MOD;
segAdd(1, dfsIndex[x], dfsIndex[x] + subtreeSize[x] - 1, z);
}
else if (op == 4) {
scanf("%d", &x);
printf("%d\n", segQuery(1, dfsIndex[x], dfsIndex[x] + subtreeSize[x] - 1));
}
}
return 0;
}
위 코드는 BOJ(백준 온라인 저지) 혹은洛谷의 트리 체인 분할 기본 문제를 해결할 수 있는 구현체이다. 실제 문제에 맞게 입력 형식과 모듈러 값을 조정하여 사용하면 된다.