TIL: 세그먼트 트리와 지속적인 배열을 이용한 고급 쿼리 처리

문제 설명

비어있는 배열 A가 주어졌을 때, 다음과 같은 5가지 연산을 지원하는 프로그램을 작성해야 한다.

1. 1 x : 배열 A의 끝에 x를 추가한다.
2. 2 L R x : A[L] ~ A[R] 범위에서 x와 xor 연산을 했을 때 최대 값이 되는 y를 찾아 출력한다.
3. 3 k : 배열 A의 마지막 k개의 원소를 제거한다.
4. 4 L R x : A[L] ~ A[R] 범위에서 x 이하의 원소 개수를 출력한다.
5. 5 L R k : A[L] ~ A[R] 범위에서 k번째로 작은 원소를 출력한다.

제한 조건이 상당히 크므로, 이진 검색 트리(Persistent Segment Tree)를 활용하여 빠르게 쿼리를 해결해야 한다.

해결 방법

데이터 구조 선택

이 문제는 배열의 특정 범위를 빠르게 쿼리해야 하기 때문에 세그먼트 트리를 활용한다.
또한, 배열이 지속적으로 변화하며 이전 상태도 필요하므로 영속적인 세그먼트 트리(Persistent Segment Tree) 를 사용한다.

Persistent Segment Tree(영속적 세그먼트 트리)란?

• 트리의 각 노드는 변경이 발생할 때 새로운 버전을 생성하여 이전 버전도 유지한다.
• O(logN)의 시간 복잡도로 특정 범위의 데이터 쿼리가 가능하다.
• 스냅샷을 저장하는 방식과 달리 최소한의 메모리를 사용한다.

코드

#include <bits/stdc++.h>
using namespace std;
const int MAX_SIZE = 1 << 19;

// 세그먼트 트리를 위한 노드 구조체
struct SegmentNode {
    int left, right, value;
    SegmentNode() : left(0), right(0), value(0) {}
    SegmentNode(int left_, int right_, int value_) : left(left_), right(right_), value(value_) {}
};

int elementCount, queryCount;
vector<SegmentNode> segmentTree(1 << 20);
int treeRoots[500005];

// 초기화: 트리를 기본 구조로 설정
void initializeTree() {
    for (int i = 1; i < MAX_SIZE; ++i) {
        segmentTree[i].left = i << 1;
        segmentTree[i].right = i << 1 | 1;
    }
}

// 세그먼트 트리 업데이트 함수
void updateTree(int node, int start, int end, int value, int index) {
    segmentTree[node].value += value;
    if (start != end) {  // 리프 노드가 아닌 경우
        int mid = (start + end) >> 1;
        int leftChild = segmentTree[node].left, rightChild = segmentTree[node].right;
        if (index <= mid) {
            segmentTree[node].left = segmentTree.size();
            segmentTree.push_back(segmentTree[leftChild]);
            updateTree(segmentTree[node].left, start, mid, value, index);
        }
        else {
            segmentTree[node].right = segmentTree.size();
            segmentTree.push_back(segmentTree[rightChild]);
            updateTree(segmentTree[node].right, mid + 1, end, value, index);
        }
    }
}

// XOR 연산을 활용한 최대 쿼리
int xorQuery(int start, int end, int x) {
    int leftBound = 0, rightBound = MAX_SIZE - 1;
    int bitShift = 18;
    int result = 0;
    while (leftBound != rightBound) {
        int mid = (leftBound + rightBound) >> 1;
        int leftSize = segmentTree[segmentTree[end].left].value - segmentTree[segmentTree[start].left].value;
        int rightSize = segmentTree[segmentTree[end].right].value - segmentTree[segmentTree[start].right].value;
        if ((x & (1 << bitShift)) && leftSize || !rightSize) {
            start = segmentTree[start].left; 
            end = segmentTree[end].left;
            rightBound = mid;
        }
        else {
            result |= (1 << bitShift);
            start = segmentTree[start].right; 
            end = segmentTree[end].right;
            leftBound = mid + 1;
        }
        --bitShift;
    }
    return result;
}

// 특정 값 이하의 개수 구하기
int countLessThanOrEqual(int start, int end, int x) {
    int leftBound = 0, rightBound = MAX_SIZE - 1;
    int count = 0;
    while (leftBound != rightBound) {
        int mid = (leftBound + rightBound) >> 1;
        if (x <= mid) {
            start = segmentTree[start].left; 
            end = segmentTree[end].left;
            rightBound = mid;
        }
        else {
            count += segmentTree[segmentTree[end].left].value - segmentTree[segmentTree[start].left].value;
            end = segmentTree[end].right; 
            start = segmentTree[start].right;
            leftBound = mid + 1;
        }
    }
    return count + (segmentTree[end].value - segmentTree[start].value);
}

// k번째 작은 원소 찾기
int findKthSmallest(int start, int end, int k) {
    int leftBound = 0, rightBound = MAX_SIZE - 1;
    while (leftBound != rightBound) {
        int mid = (leftBound + rightBound) >> 1;
        int leftSize = segmentTree[segmentTree[end].left].value - segmentTree[segmentTree[start].left].value;
        if (leftSize >= k) {
            start = segmentTree[start].left; 
            end = segmentTree[end].left;
            rightBound = mid;
        }
        else {
            start = segmentTree[start].right; 
            end = segmentTree[end].right;
            k -= leftSize;
            leftBound = mid + 1;
        }
    }
    return leftBound;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> queryCount;
    initializeTree(); 
    treeRoots[0] = 1;
    for (int i = 0; i < queryCount; ++i) {
        int queryType; cin >> queryType;
        if (queryType == 1) {
            int value; cin >> value;
            treeRoots[++elementCount] = segmentTree.size();
            segmentTree.push_back(segmentTree[treeRoots[elementCount - 1]]);
            updateTree(treeRoots[elementCount], 0, MAX_SIZE - 1, 1, value);
        }
        else if (queryType == 2) {
            int left, right, x; cin >> left >> right >> x;
            cout << xorQuery(treeRoots[left - 1], treeRoots[right], x) << '\n';
        }
        else if (queryType == 3) {
            int x; cin >> x;
            elementCount -= x;
        }
        else if (queryType == 4) {
            int left, right, x; cin >> left >> right >> x;
            cout << countLessThanOrEqual(treeRoots[left - 1], treeRoots[right], x) << '\n';
        }
        else if (queryType == 5) {
            int left, right, k; cin >> left >> right >> k;
            cout << findKthSmallest(treeRoots[left - 1], treeRoots[right], k) << '\n';
        }
    }
    return 0;
}

So...

이 문제는 Persistent Segment Tree를 사용하여 동적인 배열을 효율적으로 관리하는 방법을 학습할 수 있었다.
특히, XOR 최대값 쿼리, K번째 최소값 쿼리, 구간 내 개수 쿼리 등 여러 복잡한 연산을 효율적으로 해결해야 했다.

이번 문제를 통해 트리 기반 데이터 구조의 효율적인 활용법과 이진 탐색을 응용한 쿼리 처리 기법을 깊이 이해할 수 있었다.