[c++] 세그먼트 트리(Segment Tree) 데이터 구조

08 Dec 2023

세그먼트 트리는 배열이나 리스트와 같은 자료 구조에서 구간 쿼리를 빠르게 처리할 수 있도록 도와주는 자료 구조입니다. 세그먼트 트리는 공간 복잡도 O(n)과 시간 복잡도 O(log n)을 가지며, 주로 구간 합, 최댓값, 최솟값 등을 빠르게 계산할 때 사용됩니다.

구현

#include <vector>

using namespace std;

class SegmentTree {
private:
    vector<int> tree;
    int n;

public:
    SegmentTree(vector<int>& nums) {
        n = nums.size();
        tree.resize(4 * n);
        build(nums, 1, 0, n - 1);
    }

    void build(vector<int>& nums, int node, int start, int end) {
        if (start == end) {
            tree[node] = nums[start];
            return;
        }

        int mid = (start + end) / 2;
        build(nums, node * 2, start, mid);
        build(nums, node * 2 + 1, mid + 1, end);
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

    void update(int index, int val) {
        update(1, 0, n - 1, index, val);
    }

    void update(int node, int start, int end, int index, int val) {
        if (start == end) {
            tree[node] = val;
            return;
        }

        int mid = (start + end) / 2;
        if (index <= mid) {
            update(node * 2, start, mid, index, val);
        } else {
            update(node * 2 + 1, mid + 1, end, index, val);
        }
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

    int query(int left, int right) {
        return query(1, 0, n - 1, left, right);
    }

    int query(int node, int start, int end, int left, int right) {
        if (right < start || end < left) {
            return 0;
        }

        if (left <= start && end <= right) {
            return tree[node];
        }

        int mid = (start + end) / 2;
        return query(node * 2, start, mid, left, right) + query(node * 2 + 1, mid + 1, end, left, right);
    }
};

사용 예제

#include <iostream>
#include <vector>

using namespace std;

int main() {
    vector<int> nums = {1, 3, 5, 7, 9, 11};
    SegmentTree st(nums);

    cout << st.query(1, 3) << endl;  // 15
    st.update(1, 6);
    cout << st.query(1, 3) << endl;  // 18

    return 0;
}

세그먼트 트리는 구간 쿼리를 빠르게 처리할 수 있으며, 위의 예제에서 보듯이 빠른 시간 내에 구간의 합을 구할 수 있습니다.

더 많은 정보를 알고 싶으시다면, 아래의 참고 자료를 참고하시기 바랍니다.