세그먼트 트리와 RMQ

세그먼트 트리는 부분합 또는 RMQ(Range Minimum Query)에 적합한 자료구조다. 부분합 구하기는 펜윅 트리를 써도 되기 때문에 RMQ만 알아보겠다.

RMQ를 위한 세그먼트 트리는 원소 ($ A[1..n] $) 들에 대해 다음 두 연산을 ($ O(log \, n) $) 으로 효율적으로 처리한다.

1. query(l, r): ($ A[l..r] $) 내에서 최솟값을 반환한다.

2. update(i, x): ($ A[i] = x $) 로 갱신한다.

세그먼트 트리에는 이 외에도 구간 갱신(range update) & 단일 원소 질의(single element query), 지연 전파(lazy propagation) 등의 개념이 있다. 세그먼트 트리의 기본과 이러한 고급 개념들은 다음 글들에 자세히 설명되어 있다.

• https://www.acmicpc.net/blog/view/9
• https://www.acmicpc.net/blog/view/26
• http://codeforces.com/blog/entry/18051
• http://codeforces.com/blog/entry/15890

그리고 HackerRank에 단순히 RMQ를 구현하는 연습 문제가 있다.

RMQ를 활용하는 문제를 풀어보자.

문제: 굉장한 학생

입력이 시험 등수이기 때문에 숫자가 작을 수록 성적이 좋은 것에 주의해야 한다. 펜윅 트리 문제를 풀 때는 전수 탐색하면 ($ O(n^2) $)인 것을 ($ O(n log n) $)으로 줄였는데 이번에는 전수 탐색이 ($ O(n^3) $) 이다. 이대로는 RMQ를 어떻게 써야 할지 감도 안 오니 일단 전수 탐색을 ($ O(n^2) $)으로 줄여보자.

각 학생의 시험 점수를 3짝 ($ (A_i, B_i, C_i) $)으로 묶고 ($ A_i $) 즉 첫 번째 시험 등수를 기준으로 오름차순 정렬한다. 그러면 첫 시험을 잘 본 학생부터 앞에 온다. 학생들을 앞에서부터 1번 ~ n번으로 번호를 매긴다.

($ i < j $)일 때 항상 ($ A_i < A_j $) 이므로 ($ j $)번 학생이 굉장한 학생일 조건은 다음과 같다.

모든 ($ i < j $)에 대해 ($ B_i < B_j \; AND \; C_i < C_j $)가 성립하지 않는다.

이제 ($ j $)번 학생에 대한 ($ query(1, B_j) $)을 다음과 같이 정의한다.

($ query(1, B_j) $) = 시험 등수가 ($ B_i < B_j \; (i < j) $)인 ($ i $)번 학생들 중 최소인 ($ C_i $)

만약 ($ C_i > C_j $)라면 ($ j $)번 학생보다 대단한 학생이 한 명도 없는 것이다. 따라서 ($ j $)번 학생은 굉장한 학생이다.

질의를 한 후에는 트리의 ($ B_j $)번 노드의 값을 ($ C_j $) 로 갱신한다. 이 과정을 다음과 같이 표현할 수 있다.

   int answer = 0;

    RMQ rmq(1, n);

    for(int i=0; i<n; i++){

        int snd = scores[i].second;

        int trd = scores[i].third;

        if(rmq.query(1,snd) > trd) answer += 1;

        rmq.update(snd, trd);

    }

다음은 완전한 소스 코드다.

// ID: 2336 (굉장한 학생)

// URL: https://www.acmicpc.net/problem/2336

// Category: Segment Tree (RMQ)

#include <stdio.h>

#include <string.h>

#include <algorithm>

using namespace std;

const int INF = 1e9;

const int MAX = 500001;

struct Score { int first, second, third; } scores[MAX];

bool compareScore(const Score& a, const Score& b) {

    return a.first < b.first;

}

// segment tree - range minimum query

struct RMQ {

    int start, end;

    RMQ *left, *right;

    int minimum;

    RMQ(int l, int r) {

        start = l; end = r; minimum = INF;

        if(l == r){

            left = right = nullptr;

        }else{

            int mid = (l + r) / 2;

            left = new RMQ(l, mid);

            right = new RMQ(mid+1, r);

        }

    }

    void update(int i, int x) {

        if(start <= i && i <= end){

            if(left || right){

                left->update(i, x);

                right->update(i, x);

                minimum = min(left->minimum, right->minimum);

            }else{

                minimum = x;

            }

        }

    }

    int query(int l, int r) {

        if(r < start || end < l) return INF;

        if(l <= start && end <= r) return minimum;

        return min(left->query(l, r), right->query(l, r));

    }

};

int main() {

    int n;

    scanf("%d", &n);

    for(int i=1; i<=n; i++){

        int x; scanf("%d", &x);

        scores[x-1].first = i;

    }

    for(int i=1; i<=n; i++){

        int x; scanf("%d", &x);

        scores[x-1].second = i;

    }

    for(int i=1; i<=n; i++){

        int x; scanf("%d", &x);

        scores[x-1].third = i;

    }

    sort(scores, scores + n, compareScore);

    int answer = 0;

    RMQ rmq(1, n);

    for(int i=0; i<n; i++){

        int snd = scores[i].second;

        int trd = scores[i].third;

        if(rmq.query(1,snd) > trd) answer += 1;

        rmq.update(snd, trd);

    }

    printf("%d\n", answer);

    return 0;

}

위에 링크한 글들과는 다르게 세그먼트 트리를 진짜 트리 구조로 표현했는데, 메모리와 속도 면에서 모두 안 좋은 구현이다. 하지만 나에겐 이게 더 이해하고 사용하기 쉽고, 메모리나 시간 제한이 정말 촉박한 경우를 제외하고 지금까지 성능 면에서 큰 문제는 없었다.

코드포스에 다양한 세그먼트 트리 문제를 모아놓은 글이 있다.