TIL: Heavy-Light Decomposition을 활용한 트리 경로 연산 최적화

문제 설명

메시는 금고를 트리 구조로 관리하며, 다양한 트리 쿼리 연산을 수행해야 한다.

가능한 연산:
1. 1 X V: 금고 X의 서브트리에 있는 모든 금고에 V원을 더한다.
2. 2 X Y V: 금고 X에서 Y까지의 경로에 있는 모든 금고에 V원을 더한다.
3. 3 X V: 금고 X의 서브트리에 있는 모든 금고의 돈을 V배 한다.
4. 4 X Y V: 금고 X에서 Y까지의 경로에 있는 모든 금고의 돈을 V배 한다.
5. 5 X: 금고 X의 서브트리에 있는 모든 금고의 돈 합을 출력한다.
6. 6 X Y: 금고 X에서 Y까지의 경로에 있는 모든 금고의 돈 합을 출력한다.

모든 연산을 효율적으로 수행하기 위해 Heavy-Light Decomposition (HLD) 를 활용한다.

해결 방법

1. Heavy-Light Decomposition (HLD)

• 트리에서 경로 연산을 O(log N) 으로 수행하기 위해 사용.
• 트리를 Heavy Edge 와 Light Edge 로 분할하여 세그먼트 트리와 함께 사용.
• 경로 질의(X-Y)는 Heavy Chain을 따라가면서 세그먼트 트리에 질의.

2. 세그먼트 트리를 활용한 Lazy Propagation

• 서브트리 업데이트 (X의 모든 하위 노드 갱신)
  • entryIndex[X] ~ exitIndex[X] 범위에 대해 세그먼트 트리 연산 수행.
• 경로 업데이트 (X에서 Y까지 연산)
  • HLD를 통해 X-Y 경로를 분할하여 세그먼트 트리 연산 수행.

코드

#include<bits/stdc++.h>
#define RANGE int node, int start, int end
#define ll unsigned int
#define MAX_NODES 500001
using namespace std;

vector<int> treeGraph[MAX_NODES], heavyLightTree[MAX_NODES];
ll segmentTree[1 << 20], lazy[1 << 20][2];
int chainHead[MAX_NODES], entryIndex[MAX_NODES], exitIndex[MAX_NODES];
int subtreeSize[MAX_NODES], parent[MAX_NODES], depth[MAX_NODES];
int totalNodes, totalQueries, nodeCounter;

void initialize() {
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> totalNodes >> totalQueries;
    for (int i = 0; i < totalNodes - 1; i++) {
        int u, v; cin >> u >> v;
        treeGraph[u].push_back(v);
        treeGraph[v].push_back(u);
    }
}

// DFS를 이용한 서브트리 크기 계산
void dfsSubtreeSize(int node) {
    subtreeSize[node] = 1;
    for (int &child : treeGraph[node]) {
        if (!subtreeSize[child]) {
            depth[child] = depth[node] + 1;
            parent[child] = node;
            dfsSubtreeSize(child);
            subtreeSize[node] += subtreeSize[child];
            heavyLightTree[node].push_back(child);
            if (subtreeSize[child] > subtreeSize[heavyLightTree[node][0]]) 
                swap(heavyLightTree[node][0], heavyLightTree[node].back());
        }
    }
}

// DFS를 이용한 HLD 적용
void dfsHLD(int node) {
    entryIndex[node] = ++nodeCounter;
    for (int &child : heavyLightTree[node]) {
        chainHead[child] = (child == heavyLightTree[node][0]) ? chainHead[node] : child;
        dfsHLD(child);
    }
    exitIndex[node] = nodeCounter;
}

// Lazy Propagation을 이용한 세그먼트 트리 갱신
void propagateLazy(RANGE) {
    if (lazy[node][0] == 1 && !lazy[node][1]) return;
    segmentTree[node] *= lazy[node][0];
    segmentTree[node] += (end - start + 1) * lazy[node][1];
    if (start != end) {
        lazy[node << 1][0] *= lazy[node][0];
        lazy[node << 1][1] *= lazy[node][0];
        lazy[node << 1][1] += lazy[node][1];

        lazy[node << 1 | 1][0] *= lazy[node][0];
        lazy[node << 1 | 1][1] *= lazy[node][0];
        lazy[node << 1 | 1][1] += lazy[node][1];
    }
    lazy[node][0] = 1, lazy[node][1] = 0;
}

// 세그먼트 트리 업데이트
ll updateSegmentTree(RANGE, int left, int right, ll multiplier, ll increment) {
    propagateLazy(node, start, end);
    if (start > right || end < left) return segmentTree[node];
    if (start >= left && end <= right) {
        lazy[node][0] = multiplier;
        lazy[node][1] = increment;
        propagateLazy(node, start, end);
        return segmentTree[node];
    }
    int mid = (start + end) >> 1;
    return segmentTree[node] = updateSegmentTree(node << 1, start, mid, left, right, multiplier, increment) +
                               updateSegmentTree(node << 1 | 1, mid + 1, end, left, right, multiplier, increment);
}

// 세그먼트 트리 질의
ll querySegmentTree(RANGE, int left, int right) {
    propagateLazy(node, start, end);
    if (start > right || end < left) return 0;
    if (start >= left && end <= right) return segmentTree[node];
    int mid = (start + end) >> 1;
    return querySegmentTree(node << 1, start, mid, left, right) +
           querySegmentTree(node << 1 | 1, mid + 1, end, left, right);
}

// 경로 업데이트 (HLD)
void updateHLD(int u, int v, int multiplier, int increment) {
    while (chainHead[u] != chainHead[v]) {
        if (depth[chainHead[u]] < depth[chainHead[v]]) swap(u, v);
        updateSegmentTree(1, 1, totalNodes, entryIndex[chainHead[u]], entryIndex[u], multiplier, increment);
        u = parent[chainHead[u]];
    }
    if (depth[u] > depth[v]) swap(u, v);
    updateSegmentTree(1, 1, totalNodes, entryIndex[u], entryIndex[v], multiplier, increment);
}

// 경로 질의 (HLD)
ll queryHLD(int u, int v) {
    ll result = 0;
    while (chainHead[u] != chainHead[v]) {
        if (depth[chainHead[u]] < depth[chainHead[v]]) swap(u, v);
        result += querySegmentTree(1, 1, totalNodes, entryIndex[chainHead[u]], entryIndex[u]);
        u = parent[chainHead[u]];
    }
    if (depth[u] > depth[v]) swap(u, v);
    return result + querySegmentTree(1, 1, totalNodes, entryIndex[u], entryIndex[v]);
}

// 명령어 처리
void processQueries() {
    while (totalQueries--) {
        int type, u, v, multiplier, increment;
        cin >> type >> u;
        if (type == 1) {
            cin >> increment;
            updateSegmentTree(1, 1, totalNodes, entryIndex[u], exitIndex[u], 1, increment);
        } else if (type == 2) {
            cin >> v >> increment;
            updateHLD(u, v, 1, increment);
        } else if (type == 3) {
            cin >> multiplier;
            updateSegmentTree(1, 1, totalNodes, entryIndex[u], exitIndex[u], multiplier, 0);
        } else if (type == 4) {
            cin >> v >> multiplier;
            updateHLD(u, v, multiplier, 0);
        } else if (type == 5) {
            cout << querySegmentTree(1, 1, totalNodes, entryIndex[u], exitIndex[u]) << '\n';
        } else {
            cin >> v;
            cout << queryHLD(u, v) << '\n';
        }
    }
}

int main() {
    initialize();
    dfsSubtreeSize(1);
    dfsHLD(1);
    processQueries();
}

So...

• HLD를 활용한 경로 연산 최적화
• Lazy Propagation을 이용한 세그먼트 트리 업데이트
• 모든 연산을 ( O(\log N) ) 에 수행 가능

이 문제를 통해 HLD와 세그먼트 트리의 결합을 최적으로 활용하는 방법을 익힐 수 있었다.