세그먼트 트리를 활용한 간선 생성 최적화 기법

문제 도입

그래프 이론 문제를 해결하다 보면, 특정 노드가 구간 내의 모든 노드와 간선을 연결해야 하는 상황이 종종 발생한다. 예를 들어, "노드 u에서 구간 [L, R]에 속하는 모든 노드로 이동 가능"과 같은 제약 조건이 대표적인 경우다. 이러한 경우 단순히 모든 노드를 순회하며 간선을 생성하면 O(n²)의 시간 복잡도가 발생하여 시간 초과가 발생할 수 있다.

이러한 문제를 해결하기 위한 효율적인 기법이 바로 세그먼트 트리를 활용한 간선 생성 최적화다. 세그먼트 트리의 각 노드는 특정 구간을 나타내므로, 점과 구간 사이의 관계를 점과 세그먼트 트리 노드 사이의 관계로 변환할 수 있다.

다음으로는 이 기법을 적용한 구체적인 문제를 통해 동작 원리를 살펴보도록 하자.

예제 문제 1: Legacy (CodeForces - 787D)

문제 설명

n개의 행성이 있는 우주에서 현재 위치 s번 행성에서 다른 행성으로 이동하는 q가지 방식이 주어진다. 이동 유형은 다음 세 가지로 구분된다:

  • 유형 1: 행성 u에서 행성 v로 이동, 비용 w
  • 유형 2: 행성 u에서 구간 [L, R] 내의 임의의 행성으로 이동, 비용 w
  • 유형 3: 구간 [L, R] 내의 임의의 행성에서 행성 u로 이동, 비용 w

행성 s에서 각 행성까지의 최소 비용을 구하고, 도달 불가능한 경우 -1을 출력한다.

해결 전략

두 개의 세그먼트 트리를 사용하여 간선을 구축한다. 첫 번째 세그먼트 트리(A)는 하향식 구조로, 두 번째 세그먼트 트리(B)는 상향식 구조로 구성한다. 각 세그먼트 트리 노드는 고유 번호를 할당받으며, 이 번호를 그래프의 정점으로 활용한다.

  • 유형 1: 기존 방법과 동일하게 직접 간선을 연결
  • 유형 2: 노드 u와 세그먼트 트리 B의 구간 [L, R]에 해당하는 노드들을 연결
  • 유형 3: 세그먼트 트리 A의 구간 [L, R]에 해당하는 노드들을 노드 u와 연결

구간을 나타내는 세그먼트 트리 노드들을 통해 실제 행성 노드들과 연결함으로써, 구간 내 모든 노드와의 간선 연결을 효율적으로 처리할 수 있다. 마지막으로 다익스트라 알고리즘을 적용하여 최단 경로를 구한다.

구현 코드

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<ll, int> pli;
typedef pair<int, int> pii;

const int MAXN = 100000;
const ll INFLL = 0x3f3f3f3f3f3f3f3fLL;

int nodeCnt;
int vertA[MAXN + 5], vertB[MAXN + 5];
int visited[MAXN * 15 + 5];
ll distArr[MAXN * 15 + 5];

struct Edge {
    int target;
    int weight;
};

vector<Edge> graphAdj[MAXN * 15 + 5];

struct SegNode {
    int left;
    int right;
    int id;
};

SegNode segA[5][MAXN * 4 + 5];
SegNode segB[5][MAXN * 4 + 5];

void constructTree(int idx, int l, int r, int treeType) {
    if (treeType == 1) {
        segA[treeType][idx].left = l;
        segA[treeType][idx].right = r;
        segA[treeType][idx].id = ++nodeCnt;
    } else {
        segB[treeType][idx].left = l;
        segB[treeType][idx].right = r;
        segB[treeType][idx].id = ++nodeCnt;
    }
    
    if (l == r) {
        if (treeType == 1) {
            vertA[l] = nodeCnt;
        } else {
            vertB[l] = nodeCnt;
            graphAdj[vertB[l]].push_back({vertA[l], 0});
        }
        return;
    }
    
    int mid = (l + r) >> 1;
    constructTree(idx << 1, l, mid, treeType);
    constructTree(idx << 1 | 1, mid + 1, r, treeType);
    
    if (treeType == 1) {
        graphAdj[segA[treeType][idx << 1].id].push_back({segA[treeType][idx].id, 0});
        graphAdj[segA[treeType][idx << 1 | 1].id].push_back({segA[treeType][idx].id, 0});
    } else {
        graphAdj[segB[treeType][idx].id].push_back({segB[treeType][idx << 1].id, 0});
        graphAdj[segB[treeType][idx].id].push_back({segB[treeType][idx << 1 | 1].id, 0});
    }
}

void connectRange(int idx, int l, int r, int sourceNode, int edgeWeight, int treeType) {
    if (treeType == 1) {
        if (segA[treeType][idx].left == l && segA[treeType][idx].right == r) {
            graphAdj[sourceNode].push_back({segB[treeType][idx].id, edgeWeight});
            return;
        }
    } else {
        if (segB[treeType][idx].left == l && segB[treeType][idx].right == r) {
            graphAdj[segA[treeType][idx].id].push_back({sourceNode, edgeWeight});
            return;
        }
    }
    
    if (treeType == 1) {
        int mid = (segA[treeType][idx].left + segA[treeType][idx].right) >> 1;
        if (r <= mid) connectRange(idx << 1, l, r, sourceNode, edgeWeight, treeType);
        else if (l > mid) connectRange(idx << 1 | 1, l, r, sourceNode, edgeWeight, treeType);
        else {
            connectRange(idx << 1, l, mid, sourceNode, edgeWeight, treeType);
            connectRange(idx << 1 | 1, mid + 1, r, sourceNode, edgeWeight, treeType);
        }
    } else {
        int mid = (segB[treeType][idx].left + segB[treeType][idx].right) >> 1;
        if (r <= mid) connectRange(idx << 1, l, r, sourceNode, edgeWeight, treeType);
        else if (l > mid) connectRange(idx << 1 | 1, l, r, sourceNode, edgeWeight, treeType);
        else {
            connectRange(idx << 1, l, mid, sourceNode, edgeWeight, treeType);
            connectRange(idx << 1 | 1, mid + 1, r, sourceNode, edgeWeight, treeType);
        }
    }
}

void computeShortestPath(int startNode) {
    for (int i = 1; i <= nodeCnt; ++i) {
        distArr[i] = INFLL;
    }
    distArr[vertA[startNode]] = 0;
    
    priority_queue<pli, vector<pli>, greater<pli>> pq;
    pq.push({0, vertA[startNode]});
    
    while (!pq.empty()) {
        int current = pq.top().second;
        pq.pop();
        
        if (visited[current]) continue;
        visited[current] = 1;
        
        for (const auto& edge : graphAdj[current]) {
            int nextNode = edge.target;
            if (distArr[nextNode] > distArr[current] + edge.weight) {
                distArr[nextNode] = distArr[current] + edge.weight;
                pq.push({distArr[nextNode], nextNode});
            }
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int n, q, s;
    cin >> n >> q >> s;
    
    for (int i = 1; i <= 2; ++i) {
        constructTree(1, 1, n, i);
    }
    
    while (q--) {
        int type;
        cin >> type;
        
        if (type == 1) {
            int u, v, w;
            cin >> u >> v >> w;
            graphAdj[vertA[u]].push_back({vertB[v], w});
        } else {
            type--;
            int u, l, r, w;
            cin >> u >> l >> r >> w;
            if (type == 1) {
                connectRange(1, l, r, vertA[u], w, type);
            } else {
                connectRange(1, l, r, vertB[u], w, type);
            }
        }
    }
    
    computeShortestPath(s);
    
    for (int i = 1; i <= n; ++i) {
        ll result = min(distArr[vertA[i]], distArr[vertB[i]]);
        if (result == INFLL) cout << "-1 ";
        else cout << result << " ";
    }
    cout << "\n";
    
    return 0;
}

예제 문제 2: Hash Function

문제 설명

n개의数为あり、ハッシュテーブルに 配置する。配置规则は以下の通り: a_iのハッシュ値は a_i % n で求められる。해당 위치에 이미 다른 수가 있다면空いている位置まで 뒤로 이동한다. 완성된 해시 테이블이 주어졌을 때, 원래 수를 배치한 순서를 구한다.

해결 전략

각 수 a_i가 해시 위치에 놓이지 않은 경우,当前位置를 pos라 하면, [a_i % n, pos - 1] 구간에 있는 모든 수는 현재 수보다 먼저 배치되었음이 확실하다. 이러한 선행 관계를 세그먼트 트리를 활용하여 그래프로 구성한 후, 위상 정렬을 수행한다.

위상 정렬 결과의 노드 수가 -1이 아닌 원래 원소 수보다 작으면 해시 테이블 구성이 불가능한 경우이므로 -1을 출력한다. 그렇다면 레드큐를 활용한 위상 정렬을 통해 가능한 경우 중 가장 작은 순서를 구한다.

구현 코드

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;

const int MAXN = 200000;
const ll INFLL = 0x3f3f3f3f3f3f3f3fLL;

int testCase, n, segCnt;
vector<int> resultSeq;
vector<int> graphAdj[MAXN * 10 + 5];

int arrData[MAXN + 5];
int visitedNode[MAXN * 10 + 5];
int indegree[MAXN * 10 + 5];
int leafNode[MAXN + 5];

struct SegmentTree {
    int left, right, id;
};

SegmentTree segTree[MAXN * 4 + 5];

void buildSegmentTree(int idx, int l, int r) {
    segTree[idx].left = l;
    segTree[idx].right = r;
    segTree[idx].id = ++segCnt;
    indegree[segTree[idx].id] = 2;
    
    if (l == r) {
        indegree[segCnt] = 0;
        leafNode[l] = segCnt;
        visitedNode[segCnt] = l;
        return;
    }
    
    int mid = (l + r) >> 1;
    buildSegmentTree(idx << 1, l, mid);
    buildSegmentTree(idx << 1 | 1, mid + 1, r);
    
    graphAdj[segTree[idx << 1].id].push_back(segTree[idx].id);
    graphAdj[segTree[idx << 1 | 1].id].push_back(segTree[idx].id);
}

void addDependency(int idx, int l, int r, int targetNode) {
    if (segTree[idx].left == l && segTree[idx].right == r) {
        graphAdj[segTree[idx].id].push_back(targetNode);
        indegree[targetNode]++;
        return;
    }
    
    int mid = (segTree[idx].left + segTree[idx].right) >> 1;
    if (r <= mid) {
        addDependency(idx << 1, l, r, targetNode);
    } else if (l > mid) {
        addDependency(idx << 1 | 1, l, r, targetNode);
    } else {
        addDependency(idx << 1, l, mid, targetNode);
        addDependency(idx << 1 | 1, mid + 1, r, targetNode);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> testCase;
    
    while (testCase--) {
        cin >> n;
        
        int validCount = 0;
        for (int i = 1; i <= n; ++i) {
            cin >> arrData[i];
            if (arrData[i] != -1) {
                validCount++;
            }
        }
        
        if (validCount == 0) {
            cout << "\n";
            continue;
        }
        
        segCnt = 0;
        resultSeq.clear();
        buildSegmentTree(1, 1, n);
        
        priority_queue<pii, vector<pii>, greater<pii>> minHeap;
        
        for (int i = 1; i <= n; ++i) {
            if (arrData[i] == -1) continue;
            
            int hashPos = arrData[i] % n;
            if (hashPos == i - 1) {
                minHeap.push({arrData[i], leafNode[i]});
            } else if (hashPos > i - 1) {
                addDependency(1, hashPos + 1, n, leafNode[i]);
                if (i >= 2) {
                    addDependency(1, 1, i - 1, leafNode[i]);
                }
            } else {
                if (hashPos + 1 <= i - 1) {
                    addDependency(1, hashPos + 1, i - 1, leafNode[i]);
                }
            }
        }
        
        while (!minHeap.empty()) {
            int currentNode = minHeap.top().second;
            minHeap.pop();
            
            if (visitedNode[currentNode] != 0) {
                resultSeq.push_back(visitedNode[currentNode]);
            }
            
            for (int nextNode : graphAdj[currentNode]) {
                if (--indegree[nextNode] == 0) {
                    if (visitedNode[nextNode]) {
                        minHeap.push({arrData[visitedNode[nextNode]], nextNode});
                    } else {
                        minHeap.push({0, nextNode});
                    }
                }
            }
        }
        
        if (resultSeq.size() == validCount) {
            for (size_t i = 0; i < resultSeq.size(); ++i) {
                cout << arrData[resultSeq[i]];
                if (i == resultSeq.size() - 1) cout << "\n";
                else cout << " ";
            }
        } else {
            cout << "-1\n";
        }
        
        for (int i = 0; i <= segCnt; ++i) {
            visitedNode[i] = indegree[i] = 0;
            graphAdj[i].clear();
        }
    }
    
    return 0;
}

결론

세그먼트 트리를 활용한 간선 생성 최적화는 점과 구간 사이의 관계를 효과적으로 처리할 수 있게 해준다. 이 기법의 핵심은 구간을 세그먼트 트리의 노드로 변환하여, 구간 내 모든 원소와의 개별 연결 대신 트리 노드를 통한 간접 연결을 사용하는 것이다. 이를 통해 시간 복잡도를 크게 절감할 수 있으며, 특히 최단 경로 문제나 위상 정렬 문제에서 강력한 도구로 활용될 수 있다.

태그: segment-tree graph-algorithm edge-construction Dijkstra topological-sort

7월 13일 21:20에 게시됨