이 문제는 트리 위에서 두 가지 쿼리를 처리해야 하는 고전적인 HLD(Heavy-Light Decomposition) 적용 문제입니다. 노드 구간의 지연 갱신과 서브트리 합 질의를 효율적으로 해결해야 합니다.
핵심 개념
트리를 선형 구조로 변환하여 구간 자료구조를 적용하는 것이 핵심입니다. DFS 순서를 활용하면 서브트리를 연속된 구간으로 표현할 수 있으며, Heavy-Light Decomposition을 통해 임의의 경로를 O(log²N) 개의 구간으로 분해할 수 있습니다.
첫 번째 DFS: 트리 정보 수집
각 노드의 깊이, 부모, 서브트리 크기, 그리고 heavy child(가장 큰 서브트리를 가진 자식)를 계산합니다.
두 번째 DFS: 체인 구성 및 인덱싱
각 노드에 체인의 최상위 노드(top)와 DFS 순서번호(pos)를 할당합니다. 같은 체인 내에서는 연속된 인덱스를 가지게 됩니다.
세그먼트 트리 연산
지연 전파(lazy propagation)를 활용한 구간 신과 구간 합 질의를 지원해야 합니다.
경로 갱신 (u에서 v까지)
void addOnPath(int u, int v, long long delta) {
while (chainTop[u] != chainTop[v]) {
if (depth[chainTop[u]] < depth[chainTop[v]])
swap(u, v);
// u가 속한 체인의 구간 [pos[chainTop[u]], pos[u]] 갱신
segTree.update(pos[chainTop[u]], pos[u], delta);
u = parent[chainTop[u]];
}
// 같은 체인 내에서 처리
if (pos[u] > pos[v]) swap(u, v);
segTree.update(pos[u], pos[v], delta);
}
서브트리 합 질의
DFS 순서의 연속성을 이용하면 간단히 해결됩니다. 노드 x를 루트로 하는 서브트리는 구간 [pos[x], pos[x] + subtreeSize[x] - 1]에 대응됩니다.
long long querySubtree(int x) {
return segTree.query(pos[x], pos[x] + subtreeSize[x] - 1);
}
전체 구현
#include <bits/stdc++.h>
using namespace std;
struct SegmentTree {
struct Node {
long long sum = 0, lazy = 0;
};
vector<Node> tree;
int n;
SegmentTree(int size = 0) { init(size); }
void init(int size) {
n = size;
tree.assign(n * 4, Node());
}
void push(int idx, int l, int r) {
if (tree[idx].lazy == 0) return;
long long val = tree[idx].lazy;
tree[idx].sum += val * (r - l + 1);
if (l != r) {
tree[idx*2].lazy += val;
tree[idx*2+1].lazy += val;
}
tree[idx].lazy = 0;
}
void pull(int idx) {
tree[idx].sum = tree[idx*2].sum + tree[idx*2+1].sum;
}
void rangeAdd(int idx, int l, int r, int ql, int qr, long long val) {
push(idx, l, r);
if (qr < l || r < ql) return;
if (ql <= l && r <= qr) {
tree[idx].lazy += val;
push(idx, l, r);
return;
}
int mid = (l + r) >> 1;
rangeAdd(idx*2, l, mid, ql, qr, val);
rangeAdd(idx*2+1, mid+1, r, ql, qr, val);
pull(idx);
}
long long rangeQuery(int idx, int l, int r, int ql, int qr) {
push(idx, l, r);
if (qr < l || r < ql) return 0;
if (ql <= l && r <= qr) return tree[idx].sum;
int mid = (l + r) >> 1;
return rangeQuery(idx*2, l, mid, ql, qr)
+ rangeQuery(idx*2+1, mid+1, r, ql, qr);
}
};
struct Edge {
int to, next;
};
class HeavyLight {
public:
vector<int> head, depth, parent, heavy, top, pos, sz;
vector<Edge> edges;
int n, curPos, edgeCnt;
SegmentTree seg;
HeavyLight(int size) {
n = size;
head.assign(n + 1, -1);
depth.resize(n + 1);
parent.resize(n + 1);
heavy.assign(n + 1, -1);
top.resize(n + 1);
pos.resize(n + 1);
sz.resize(n + 1);
edges.reserve(n * 2);
edgeCnt = 0;
seg.init(n);
}
void addEdge(int u, int v) {
edges.push_back({v, head[u]});
head[u] = edgeCnt++;
edges.push_back({u, head[v]});
head[v] = edgeCnt++;
}
int dfs(int u, int p) {
parent[u] = p;
sz[u] = 1;
int maxSub = 0;
for (int e = head[u]; ~e; e = edges[e].next) {
int v = edges[e].to;
if (v == p) continue;
depth[v] = depth[u] + 1;
int sub = dfs(v, u);
if (sub > maxSub) {
maxSub = sub;
heavy[u] = v;
}
sz[u] += sub;
}
return sz[u];
}
void decompose(int u, int t, int& idx) {
top[u] = t;
pos[u] = ++idx;
if (heavy[u] != -1) {
decompose(heavy[u], t, idx);
}
for (int e = head[u]; ~e; e = edges[e].next) {
int v = edges[e].to;
if (v == parent[u] || v == heavy[u]) continue;
decompose(v, v, idx);
}
}
void pathUpdate(int u, int v, long long val) {
while (top[u] != top[v]) {
if (depth[top[u]] < depth[top[v]]) swap(u, v);
seg.rangeAdd(1, 1, n, pos[top[u]], pos[u], val);
u = parent[top[u]];
}
if (pos[u] > pos[v]) swap(u, v);
seg.rangeAdd(1, 1, n, pos[u], pos[v], val);
}
long long subtreeQuery(int u) {
return seg.rangeQuery(1, 1, n, pos[u], pos[u] + sz[u] - 1);
}
};
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, q;
if (!(cin >> n)) return 0;
HeavyLight hld(n);
for (int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
u++; v++;
hld.addEdge(u, v);
}
int idx = 0;
hld.dfs(1, 0);
hld.decompose(1, 1, idx);
cin >> q;
while (q--) {
char op;
cin >> op;
if (op == 'A') {
int u, v;
long long w;
cin >> u >> v >> w;
u++; v++;
hld.pathUpdate(u, v, w);
} else {
int u;
cin >> u;
u++;
cout << hld.subtreeQuery(u) << '\n';
}
}
return 0;
}
추가 활용: 구간 최댓값/최솟값
세그먼트 트리에 최댓값/최솟값을 저장하도록 수정하면 경로상의 최댓값/최솟값도 동일한 방식으로 구할 수 있습니다. 다만 합과 달리 결합 법칙만 성립하면 되므로 구현이 더욱 간단해집니다.
주의사항
입력 노드 번호가 0-based일 수 있으므로 적절히 1-based로 변환해야 합니다. 또한 long long 사용 여부를 신중히 판단해야 합니다.