Heavy-Light Decomposition과 세그먼트 트리를 활용한 트리 경로 최대 연속합 처리

문제 개요 및 핵심 접근법

트리 구조에서 두 정점 사이의 경로에 대해 최대 연속 부분합을 조회하고, 경로상의 모든 정점 값을 일괄 수정하는 연산을 처리해야 합니다. 정점의 개수와 쿼리의 개수가 최대 100,000에 달하므로, 단순한 순회 방식은 시간 제한을 초과합니다. 이를 효율적으로 해결하기 위해 Heavy-Light Decomposition(HLD)과 세그먼트 트리(Segment Tree)를 결합한 접근이 필요합니다. HLD는 트리 경로를 O(log N)개의 연속된 구간으로 분할하며, 세그먼트 트리는 각 구간에서 필요한 집계 정보를 빠르게 관리합니다.

세그먼트 트리 노드 설계 및 병합 로직

경로를 일차원 배열로 펼쳤을 때 최대 연속 부분합을 구하려면 각 구간에서 네 가지 정보를 유지해야 합니다. 전체 합계(sum), 왼쪽 끝에서 시작하는 최대 연속합(pref), 오른쪽 끝에서 끝나는 최대 연속합(suff), 그리고 구간 내 최대 연속합(best)입니다. 두 인접 구간을 병합할 때는 다음과 같은 논리를 적용합니다.

Node mergeNodes(const Node& left, const Node& right) {
    Node res;
    res.sum = left.sum + right.sum;
    res.pref = std::max(0, std::max(left.pref, left.sum + right.pref));
    res.suff = std::max(0, std::max(right.suff, right.sum + left.suff));
    res.best = std::max(0, std::max({left.best, right.best, left.suff + right.pref}));
    return res;
}

문제 조건에 따라 모든 값이 음수일 경우 최대 연속합은 0으로 처리되므로, 각 계산 단계에서 0과의 비교를 통해 음수 결과를 배제합니다. 빈 구간을 나타내는 초기 노드는 모든 필드가 0으로 설정되어 항등원 역할을 합니다.

경로 쿼리 시 구간 방향성 처리

HLD를 사용하면 두 정점 사이의 경로는 여러 개의 체인(Chain)으로 분할됩니다. 세그먼트 트리에서 조회한 각 체인 구간을 병합할 때, 방향성을 정확히 유지하는 것이 가장 중요합니다. LCA(최소 공통 조상)를 기준으로 왼쪽 경로와 오른쪽 경로로 나누어 생각해야 합니다.

  • 깊이가 깊은 정점에서 체인 헤드 방향으로 올라갈 때, 조회된 구간은 DFS 순서상 위에서 아래로 배치되지만 실제 경로는 아래에서 위로 향합니다.
  • 따라서 U측 경로와 V측 경로에서 수집한 구간을 별도의 누적 변수(resU, resV)로 관리합니다.
  • U측은 새 구간을 왼쪽에 병합하고, V측은 새 구간을 오른쪽에 병합하여 방향을 유지합니다.
  • 최종 병합 직전 U측 누적 결과의 pref와 suff를 교환하여 방향을 반전시킨 후, V측 결과와 순서대로 병합합니다.

Lazy Propagation 주의사항

경로 갱신 연산은 구간 대입(Assign) 형태입니다. Lazy 태그를 구현할 때 초기값을 0으로 설정하면, 실제 갱신 값이 0인 경우와 충돌할 수 있습니다. 따라서 갱신 대기 상태를 나타내기 위해 별도의 불리언 플래그(hasLazy)를 사용하거나, 문제의 값 범위를 벗어나는 센티널 값을 활용해야 합니다. 아래 구현에서는 불리언 플래그를 사용하여 0을 포함한 모든 정수 값을 안전하게 처리합니다.

전체 구현 코드

아래 코드는 HLD와 세그먼트 트리를 결합하여 경로 최대 연속합 조회 및 구간 대입 갱신을 처리하는 완전한 구현입니다. 변수명과 구조를 재구성하여 가독성을 높였으며, 방향성 병합과 Lazy 전파 로직을 명시적으로 분리했습니다.

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int MAXN = 100005;

struct Node {
    int sum, pref, suff, best;
    int lazy;
    bool hasLazy;
    Node() : sum(0), pref(0), suff(0), best(0), lazy(0), hasLazy(false) {}
};

int n, m;
int initVal[MAXN];
vector<int> adj[MAXN];
int parent[MAXN], depth[MAXN], heavy[MAXN], head[MAXN], pos[MAXN];
int curPos;
Node tree[4 * MAXN];
int mappedVal[MAXN];

// 두 구간 노드를 병합 (방향성 고려)
Node mergeNodes(const Node& l, const Node& r) {
    Node res;
    res.sum = l.sum + r.sum;
    res.pref = max(0, max(l.pref, l.sum + r.pref));
    res.suff = max(0, max(r.suff, r.sum + l.suff));
    res.best = max(0, max({l.best, r.best, l.suff + r.pref}));
    return res;
}

// Lazy 값 적용
void applyLazy(int idx, int l, int r, int val) {
    tree[idx].sum = (r - l + 1) * val;
    tree[idx].pref = tree[idx].suff = tree[idx].best = max(0, tree[idx].sum);
    tree[idx].lazy = val;
    tree[idx].hasLazy = true;
}

// Lazy 전파
void push(int idx, int l, int r) {
    if (tree[idx].hasLazy) {
        int mid = (l + r) / 2;
        applyLazy(2 * idx, l, mid, tree[idx].lazy);
        applyLazy(2 * idx + 1, mid + 1, r, tree[idx].lazy);
        tree[idx].hasLazy = false;
    }
}

// 세그먼트 트리 초기화
void build(int idx, int l, int r) {
    if (l == r) {
        int v = mappedVal[l];
        tree[idx].sum = v;
        tree[idx].pref = tree[idx].suff = tree[idx].best = max(0, v);
        return;
    }
    int mid = (l + r) / 2;
    build(2 * idx, l, mid);
    build(2 * idx + 1, mid + 1, r);
    tree[idx] = mergeNodes(tree[2 * idx], tree[2 * idx + 1]);
}

// 구간 대입 갱신
void updateRange(int idx, int l, int r, int ql, int qr, int val) {
    if (ql > r || qr < l) return;
    if (ql <= l && r <= qr) {
        applyLazy(idx, l, r, val);
        return;
    }
    push(idx, l, r);
    int mid = (l + r) / 2;
    updateRange(2 * idx, l, mid, ql, qr, val);
    updateRange(2 * idx + 1, mid + 1, r, ql, qr, val);
    tree[idx] = mergeNodes(tree[2 * idx], tree[2 * idx + 1]);
}

// 구간 조회
Node queryRange(int idx, int l, int r, int ql, int qr) {
    if (ql > r || qr < l) return Node();
    if (ql <= l && r <= qr) return tree[idx];
    push(idx, l, r);
    int mid = (l + r) / 2;
    return mergeNodes(queryRange(2 * idx, l, mid, ql, qr),
                      queryRange(2 * idx + 1, mid + 1, r, ql, qr));
}

// HLD 1단계: 서브트리 크기 계산 및 Heavy Edge 결정
int dfs(int u, int p, int d) {
    depth[u] = d;
    parent[u] = p;
    int size = 1;
    int maxSubSize = 0;
    heavy[u] = -1;
    for (int v : adj[u]) {
        if (v != p) {
            int subSize = dfs(v, u, d + 1);
            size += subSize;
            if (subSize > maxSubSize) {
                maxSubSize = subSize;
                heavy[u] = v;
            }
        }
    }
    return size;
}

// HLD 2단계: 체인 분할 및 DFS 순서 매핑
void decompose(int u, int h) {
    head[u] = h;
    pos[u] = ++curPos;
    mappedVal[curPos] = initVal[u];
    if (heavy[u] != -1) decompose(heavy[u], h);
    for (int v : adj[u]) {
        if (v != parent[u] && v != heavy[u]) decompose(v, v);
    }
}

// 경로 구간 대입 갱신
void updatePath(int u, int v, int val) {
    while (head[u] != head[v]) {
        if (depth[head[u]] < depth[head[v]]) swap(u, v);
        updateRange(1, 1, n, pos[head[u]], pos[u], val);
        u = parent[head[u]];
    }
    if (depth[u] > depth[v]) swap(u, v);
    updateRange(1, 1, n, pos[u], pos[v], val);
}

// 경로 최대 연속합 조회 (방향성 보존)
int queryPathMax(int u, int v) {
    Node resU, resV;
    while (head[u] != head[v]) {
        if (depth[head[u]] > depth[head[v]]) {
            resU = mergeNodes(queryRange(1, 1, n, pos[head[u]], pos[u]), resU);
            u = parent[head[u]];
        } else {
            resV = mergeNodes(resV, queryRange(1, 1, n, pos[head[v]], pos[v]));
            v = parent[head[v]];
        }
    }
    if (depth[u] > depth[v]) {
        resU = mergeNodes(queryRange(1, 1, n, pos[v], pos[u]), resU);
    } else {
        resV = mergeNodes(resV, queryRange(1, 1, n, pos[u], pos[v]));
    }
    // U측 경로는 아래에서 위로 수집되었으므로 좌우 최대값을 교환하여 방향 보정
    swap(resU.pref, resU.suff);
    return mergeNodes(resU, resV).best;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    if (!(cin >> n)) return 0;
    for (int i = 1; i <= n; ++i) cin >> initVal[i];
    for (int i = 0; i < n - 1; ++i) {
        int u, v; cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    curPos = 0;
    dfs(1, 0, 0);
    decompose(1, 1);
    build(1, 1, n);

    cin >> m;
    bool first = true;
    for (int i = 0; i < m; ++i) {
        int type; cin >> type;
        if (type == 1) {
            int u, v; cin >> u >> v;
            if (!first) cout << " ";
            cout << queryPathMax(u, v);
            first = false;
        } else {
            int u, v, c; cin >> u >> v >> c;
            updatePath(u, v, c);
        }
    }
    cout << "\n";
    return 0;
}

시간 복잡도 및 제약 조건

HLD를 통한 경로 분할은 최대 O(log N)개의 체인을 생성하며, 각 체인에 대한 세그먼트 트리 연산은 O(log N)이 소요됩니다. 따라서 쿼리당 시간 복잡도는 O(log² N)이며, N, M ≤ 100,000 조건에서 약 1~2초 내에 안정적으로 동작합니다. 메모리 사용량은 세그먼트 트리 배열과 그래프 인접 리스트를 포함하여 O(N) 수준으로 관리됩니다. 방향성 병합 로직과 Lazy 플래그 처리를 정확히 구현하면 엣지 케이스(음수-only 경로, 0으로의 갱신 등)에서도 안정적인 결과를 보장합니다.

태그: Heavy-Light Decomposition 세그먼트 트리 경로 쿼리 lazy propagation 최대 연속 부분합

10월 8일 04:22에 게시됨