C++ 코딩테스트 강좌, 세그먼트 트리

최근의 알고리즘 코딩 테스트에서 세그먼트 트리는 자주 등장하는 자료구조 중 하나입니다. 세그먼트 트리는 특정 구간의 합, 최소값, 최대값 등을 효율적으로 계산할 수 있도록 도와주는 자료구조입니다. 이번 글에서는 세그먼트 트리를 사용하는 알고리즘 문제를 살펴보고, 이를 해결하기 위한 과정과 실제 구현 예시를 제공하겠습니다.

문제 설명

주어진 배열에서 구간 [l, r]의 합을 구하는 쿼리를 수행할 수 있는 프로그램을 작성하시오. 입력으로는 배열의 크기 n과 쿼리의 개수 m이 주어지며, 이어서 n개의 정수가 배열로 주어집니다. 각 쿼리는 두 개의 정수 l과 r로 구성되어 있습니다.

문제 예시

문제 접근 방법

위 문제는 단순히 배열을 순회하여 합계를 계산하는 방법으로도 해결할 수 있지만, 주어진 쿼리의 수가 많고 배열의 크기가 큰 경우 비효율적입니다. 이에 따라 세그먼트 트리를 사용하면 효율적으로 문제를 해결할 수 있습니다.

세그먼트 트리란?

세그먼트 트리는 배열을 구간 단위로 나누어 각 구간의 정보를 저장하는 트리 구조입니다. 이 트리는 다음과 같은 특성을 가지고 있습니다:

각 노드는 배열의 구간을 표현합니다.
각 리프 노드는 배열의 원소를 저장합니다.
각 내부 노드는 자식 노드의 값을 합산하여 저장합니다.

세그먼트 트리를 활용하면 구간의 합을 O(log n) 시간 복잡도로 계산할 수 있습니다. 이는 배열의 원소 수가 많아질수록 그 효율성이 두드러집니다.

세그먼트 트리 구현

이제 세그먼트 트리를 이용하여 문제를 해결해보겠습니다. 다음은 C++로 세그먼트 트리를 구현하는 코드입니다.


#include <iostream>
#include <vector>
using namespace std;

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

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

    void build(int node, int start, int end) {
        if (start == end) {
            // Leaf node
            tree[node] = arr[start];
        } else {
            int mid = (start + end) / 2;
            // Recursively build the segment tree
            build(2 * node + 1, start, mid);
            build(2 * node + 2, mid + 1, end);
            // Internal node will have the sum of the two children
            tree[node] = tree[2 * node + 1] + tree[2 * node + 2];
        }
    }

    int sum(int l, int r) {
        return sum(0, 0, n - 1, l, r);
    }

    int sum(int node, int start, int end, int l, int r) {
        if (r &lt; start || end &lt; l) {
            // range represented by a node is completely outside the given range
            return 0;
        }
        if (l &lt;= start &amp;&amp; end &lt;= r) {
            // range represented by a node is completely inside the given range
            return tree[node];
        }
        // range represented by a node is partially inside and partially outside the given range
        int mid = (start + end) / 2;
        int left_child = sum(2 * node + 1, start, mid, l, r);
        int right_child = sum(2 * node + 2, mid + 1, end, l, r);
        return left_child + right_child;
    }
};

int main() {
    int n, m;
    cin &gt;&gt; n &gt;&gt; m;
    vector<int> arr(n);
    for (int i = 0; i &lt; n; ++i) {
        cin &gt;&gt; arr[i];
    }
    SegmentTree segment_tree(arr);
    
    for (int i = 0; i &lt; m; ++i) {
        int l, r;
        cin &gt;&gt; l &gt;&gt; r;
        // Adjusting for 0-based indexing
        cout &lt;&lt; segment_tree.sum(l - 1, r - 1) &lt;&lt; endl;
    }
    
    return 0;
}
</int></int></int></vector></iostream>

코드 분석

위 코드는 아래와 같은 절차로 세그먼트 트리를 구현합니다:

클래스 정의: SegmentTree 클래스를 정의하고 필요한 멤버 변수를 선언합니다. 여기에는 세그먼트 트리를 저장할 tree 벡터와 원본 배열 arr, 배열의 크기 n이 포함됩니다.
생성자: 생성자에서 입력 배열을 받아 세그먼트 트리를 초기화하고 build 메서드를 호출하여 트리를 구축합니다.
트리 구축: build 메서드는 재귀적으로 트리를 구축합니다. 리프 노드는 원본 배열의 원소를 저장하고, 내부 노드는 자식 노드의 합을 저장합니다.
합 계산: sum 메서드는 주어진 구간의 합을 계산합니다. 이 메서드는 노드가 주어진 구간의 범위에 속하는지 확인하고, 속하지 않는 경우 적절한 값을 반환하거나 재귀적으로 합을 계산합니다.
메인 함수: 사용자로부터 배열의 크기와 쿼리 수를 입력받고, 배열을 입력한 뒤 SegmentTree 객체를 생성합니다. 그리고 각 쿼리에 대해 결과를 출력합니다.

결론

세그먼트 트리는 특정 구간의 합을 빠르게 계산할 수 있도록 도와주는 강력한 도구입니다. C++를 사용하여 세그먼트 트리를 구현함으로써, 알고리즘 문제 해결에 필요한 효율성을 높일 수 있었습니다. 이 자료구조는 단순합뿐만 아니라 다양한 쿼리에도 활용할 수 있으므로, 더 많은 연습을 통해 숙달하는 것이 중요합니다.

이제 여러분은 C++에서 세그먼트 트리를 구현하고 사용할 수 있는 기반 지식을 갖추게 되었습니다. 다양한 문제를 해결하면서 이 자료구조의 힘을 확인해보세요.