경로 및 서브트리 갱신을 위한 Heavy-Light Decomposition 활용

이 문제는 트리 위에서 두 가지 쿼리를 처리해야 하는 고전적인 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 사용 여부를 신중히 판단해야 합니다.

태그: HLD Heavy-Light Decomposition segment tree lazy propagation Tree Path Query

8월 1일 09:05에 게시됨