세그먼트 트리1 (구현)

Jewook·2022년 7월 28일

알고리즘

목록 보기

8/14

RMQ란?

Range Minimum Query. 배열이 있을 때, 배열의 i~j의 최소값을 물어보는 쿼리이다. 구간에 관한 질의(query)문제의 대표격이다. 구간에 관해 다루는 알고리즘을 배우면서 처음 공부하게 되는데, 이 문제는 다양하게 응용할 수 있다. (최소값, 최대값, 구간의 합 등) 브루트포스는 당연히 쿼리마다 구간의 길이에 비례하는 시간이 걸리기 때문에 O(NQ)의 시간복잡도를 가진다. 줄일 수 있는 방법이 없을까?

Sqrt Decomposition

배열을 루트 N을 크기로 나누어 구간마다 최소값을 저장한다. 그룹의 성질을 잘 이용해서 쿼리를 sqrt N 시간복잡도만에 해결할 수 있다.

전처리 O(N)

int r = sqrt(N);
int arr[N], group[N/r+1];

for (int i=0; i < n; ++i) {
	if (i%r==0) group[i/r] = arr[i];
    else group[i/r] = min(group[i/r], arr[i]);
}

쿼리 O(sqrt N)

// i~j 구간의 최소값? 
int query(int i, int j) {
    int ret = INF; // 큰값으로 초기화
    // 같은 그룹에 속하는 경우 브루트포스
    if (i / r == j / r) {
        while (i <= j) ret = min(ret, arr[i++]);
        return ret;
    }
    // i가 속한 구간 O(sqrt N)
    while ( i % r != 0) ret = min(ret, arr[i++]);
    // j가 속한 구간 O(sqrt N)
    while ( j % r != r - 1) ret = min(ret, arr[j--]);
    // 중간에 있는 구간 O(sqrt N)
    for (int grp = i / r; grp <= j / r; ++grp) ret = min(ret, group[grp]);
    return ret;
}

Sparse Table (희소배열)

다이나믹 프로그래밍을 응용해보자.
dp[i][j] : i부터 2^j만큼 구간의 최소값이라고 할 때,
dp[i][j] = min(dp[i][j-1], dp[i+2^j-1][j-1]) 이다.

전처리 O(NlogN)

for (int i=0; i < N; ++i) dp[i][0] = arr[i];
for (int j=1; j < log N+1; ++j) {
	for (int i=0; i < N; ++i) {
    	if (i+(1<<j) > N) break;
        dp[i][j] = min(dp[i][j-1], dp[i+(1<<(j-1))][j-1]);
    }
}

쿼리 O(log N)

int query(int i, int j) {
	int ret = arr[i], k = log N;
    while ( i<=j && k>=0 ) {
    	if (i+(1<<k)-1 <= j) {
        	ret = min(ret, dp[i][k]);
            i += 1<<k;
       	}
        k--;
    }
    return ret;
}

희소배열은 정말 빠르다. 구간의 범위가 2의 거듭제곱이거나 이진수 표현상으로 2의 거듭제곱과 가깝다면 O(1)만에 쿼리를 진행할 수도 있다. 하지만 배열이 갱신되는 경우 전처리를 새로 해줘야 한다. 갱신되지만 않는다면, 정말 좋지만, 좀 더 유연하고 강력한 무기가 없을까?

세그먼트 트리

대망의,,

struct SegTree {
    int n;
    vi tree;

    //Constructor
    SegTree(const vector<int>& arr) {
        n = arr.size();
        tree.resize(n * 4); // 1<<(log2(n)+1)로 해도된다
        init(arr, 0, n - 1, 1);
    }

    // 전처리 O(NlogN)
    int init(const vector<int>& arr, int lo, int hi, int node) {
        if (lo == hi)
            return tree[node] = arr[lo];

        int mid = (lo + hi) / 2;
        int leftVal = init(arr, lo, mid, node * 2);
        int rightVal = init(arr, mid + 1, hi, node * 2 + 1);
        return tree[node] = min(leftVal, rightVal);
    }

    // Query O(logN)
    int query(int lo, int hi) {
        return query(lo, hi, 1, 0, n - 1);
    }
    int query(int lo, int hi, int node, int nodelo, int nodehi) {
        if (nodehi < lo || hi < nodelo) return INF;
        if (lo <= nodelo && nodehi <= hi) {
            return tree[node];
        }

        int mid = (nodelo + nodehi) / 2;
        int leftVal = query(lo, hi, node * 2, nodelo, mid);
        int rightVal = query(lo, hi, node * 2 + 1, mid + 1, nodehi);
        return min(leftVal, rightVal);
    }

    // Update(log N)
    void update(int idx, int newVal) {
        update(idx, newVal, 1, 0, n - 1);
    }
    void update(int idx, int newVal, int node, int nodelo, int nodehi) {
        if (idx < nodelo || nodehi < idx) return;
        if (nodelo == nodehi) {
            tree[node] = newVal;
            return;
        }

        int mid = (nodelo + nodehi) / 2;
        update(idx, newVal, node * 2 + 1, mid + 1, nodehi);
        update(idx, newVal, node * 2, nodelo, mid);
        tree[node] = min(tree[node * 2], tree[node * 2 + 1]);
    }
};

갱신까지 가능한 극강의 구간쿼리 도구이다.

팬윅트리 (Binary Indexed Tree)

struct FenwickTree{
    vector<int> tree;

    FenwickTree(int n) : tree(n+1) {}

    // 1부터이다! 0부터아님
    int sum(int idx) {
        ll ret = 0;
        while (idx > 0) {
            ret += tree[idx];
            idx &= (idx - 1); // idx -= (idx&-idx)와 같음 
            // 끝의 1비트만큼 뺌
        }
        return ret;
    }
    
    void add(int idx, ll val) {
        while (idx < tree.size()) {
            tree[idx] += val;
            idx += idx & -idx; // 끝의 1비트만큼 더함
        }
    }
};

'구간합'을 구하는 세그트리의 변종이다. 세그트리도 구간합을 구할 수 있지만, 속도와 메모리를 절약하는 구간합 한정 세그트리 최적화 버전이다.
메모리는 N만 사용하며, 속도는 O(logN)이긴 하지만 오버헤드가 싹 제거되어 매우 빠르다.