세그먼트 트리의 응용

세그먼트 트리

세그먼트 트리는 수열에 다양한 구간 질의를 빠르게 온라인으로 처리할 수 있게 하는 자료구조이다. 만약 세그먼트 트리에 대해 접해본 적이 전혀 없다면, 네 편의 강의 #1, #2, #3, #4를 참조하라. 세그먼트 트리를 설명할 때에 가장 흔하게 쓰이는 연습문제는 다음과 같다.

문제 1 수열 $A[0], \cdots, A[N -1]$ 이 있다. 이 수열에 질의를 고속으로 처리해야 한다. 질의의 종류는 아래와 같다.

두 수 $i, k$ 가 입력으로 주어진다. 이 때, $A[i]$ 를 $k$만큼 증가시켜라.
두 수 $i, j$ 가 입력으로 주어진다. 이 때, $A[i] + \cdots + A[j]$ 의 값을 구하라.

이 때 세그먼트 트리를 사용하면 각 질의를 $O(\log N)$의 시간 복잡도에 해결할 수 있다.

이는 $[1, N]$ 범위 안에서 적절한 $O(N\log N)$종류의 구간을 미리 지정해, 이 구간의 합들을 메모이제이션하는 기법이기도 하다.

임의의 $i$는 $O(\log N)$ 종류의 구간에만 포함되며,
임의의 $(i, j)$ 에 대해, 구간 $A[i..j]$는 서로 겹치지 않는 $O(\log N)$ 종류의 구간의 합 (disjoint union)으로 표현된다.

세그먼트 트리의 정점들은 단지 트리의 정점일 뿐만 아니라, 이 조건을 만족하도록 적절한 $O(N \log N)$ 개의 구간들을 고르는 문제의 절묘한 해답이기도 하다. 추가적으로, 지연 갱신 기법도 같이 사용하면 다음과 같은 문제도 해결할 수 있다.

문제 2 수열 $A[0], \cdots, A[N -1]$ 이 있다. 이 수열에 질의를 고속으로 처리해야 한다. 질의의 종류는 아래와 같다.

세 수 $i, j, k$ 가 입력으로 주어진다. 이 때, 수열의 $i$번째에서 $j$번째까지의 원소, 즉 $A[i], \cdots, A[j]$ 모두를 $k$만큼 증가시켜라.
두 수 $i, j$ 가 입력으로 주어진다. 이 때, $A[i] + \cdots + A[j]$ 의 값을 구하라.

지연 갱신 기법은, 2번 질의의 구간 $A[i..j]$가 $O(\log N)$개 정점의 구간의 합으로 표현되었듯이, 1번 질의의 구간 $A[i..j]$ 또한 $O(\log N)$개 정점의 구간의 합으로 표현될 수 있다는 사실을 이용한다. 이를 이용하면 1번 질의는 정점을 $O(\log N)$개 업데이트하는 질의가 된다. 여기에 ‘업데이트를 꼭 즉시 적용시킬 필요는 없다’라는 아이디어를 추가하면, 정점 업데이트를 더욱 효율적으로 할 수 있다. 이 정점에 적용되어야 하는 업데이트를 모두 모아놓은 뒤, 이 정점의 값을 이용해야 할 때(이 정점에 2번 쿼리가 들어왔을 때)에 한꺼번에 업데이트를 적용해주는 것이다.

위 두 문제는 세그먼트 트리의 가장 기초적인 문제들이다. 대부분의 경우, 세그먼트 트리를 배울 때에는 이들 문제를 포함해 수열에 쿼리를 적용시키는 문제들을 통해 세그먼트 트리를 익히게 된다. 아래는 세그먼트 트리의 가장 기본적인 형태이다. 지연 갱신 세그먼트 트리에는 update, query 함수를 실행할 때마다 호출될 propagate 함수가 추가되어야 한다.

struct segt {
    node t[maxn*2];
    node merge(node a, node b) { 
        // 두 노드가 합쳐진 결과 정점을 반환해야 한다.
        // 위 예제에서는 a+b에 해당한다.
    }
    void update(int s, int e, int x, int p, node v) {
        if (s==e) t[x] = v;
        else {
            int m = (s+e)/2;
            if (p<=m) update(s, m, 2*x, p, v);
            else update(m+1, e, 2*x+1, p, v);
            t[x] = merge(t[2*x], t[2*x+1]);
        }
    }
    void update(int p, node v) { return update(0, MAXN, 1, p, v); }
    node query(int s, int e, int x, int l, int r) {
        if (l<=s and e<=r) return t[x];
        else if (l<=e and s<=r) {
            int m = (s+e)/2;
            return merge(query(s, m, 2*x, l, r), query(m+1, e, 2*x+1, l, r));
        } else return I; // 다른 연산에 영향을 주지 않는 항등원 역할을 하는 정점을 반환해야 한다. 위 예제에서는 0에 해당한다.
    }
    node query(int l, int r) { return query(0, MAXN, 1, l, r); }
};

세그먼트 트리의 응용

이 글에서는 수열에 쿼리를 적용시키는 형태를 넘어서, 다양한 객체를 세그먼트 트리에 집어넣는 사례들을 살펴본다. 세그먼트 트리 연습문제들에서는 node 타입이 단순히 정수형이었지만, node 타입에 다른 다양한 구조체를 넣는 사례들을 살펴볼 것이다.

문제 1: Addition Robot

링크

이 문제는 A와 B로 이루어진 문자열과 두 종류의 쿼리가 주어지기 때문에 직관적으로 세그먼트 트리와 비슷한 형태의 자료구조를 사용해야 한다는 생각을 할 수 있다. 하지만 어떤 방법으로 업데이트를 처리해야 할 지 쉽사리 눈에 띄지 않는다. 이 때,

\[\begin{bmatrix}A+B\\B\end{bmatrix} = \begin{bmatrix}1&1\\0&1\end{bmatrix} \begin{bmatrix}A\\B\end{bmatrix}\] \[\begin{bmatrix}A\\A+B\end{bmatrix} = \begin{bmatrix}1&0\\1&1\end{bmatrix} \begin{bmatrix}A\\B\end{bmatrix}\]

라는 사실을 관찰하면, 세그먼트 트리의 각 정점에 행렬을 넣을 수 있겠다는 관찰을 할 수 있게 된다. 알파벳 A는 $\begin{bmatrix}1&1\0&1\end{bmatrix}$로, 알파벳 B는 $\begin{bmatrix}1&0\1&1\end{bmatrix}$로 생각하면, 알파벳 $S_i$에서 $S_j$까지의 연산을 차례대로 실행하는 작업은 각 알파벳에 해당하는 행렬을 차례대로 곱하는 작업으로 생각할 수 있다. 업데이트 질의는, 행렬을 토글하는 질의로 생각할 수 있다.

여기서, 지연 갱신 세그먼트 트리가 작동하는 것은 행렬곱의 결합법칙 덕분이라는 것을 관찰할 수 있다. $\begin{bmatrix}A\B\end{bmatrix}$에 행렬을 하나씩 순서대로 곱해주는 작업과, 모두 곱해진 행렬에 $\begin{bmatrix}A\B\end{bmatrix}$를 한꺼번에 곱하는 작업은 같다.

위 세그먼트 트리 예시 코드에서, node 구조체는 $2 \times 2$ 행렬, merge 함수는 행렬곱, 항등원 I는 $\begin{bmatrix}1&0\0&1\end{bmatrix}$으로 만들면 된다. 행렬의 크기를 상수로 생각한다면, 각 질의의 시간 복잡도는 $O(\log N)$이 된다.

IOI 2013 기출문제인 ‘웜뱃’도 세그먼트 트리에 행렬을 넣어 업데이트 하는 방법으로 해결할 수 있다. 다만, 이 문제에서는 행렬을 합치는 $O(C^3)$의 시간을 절약해야 하는 데다, 메모리 최적화 기법도 필요하기 때문에 난이도가 어려워진다.

문제 2: 없던 일처럼

링크

이 문제에서 ‘사건’은, 조건 분기를 만들어 전체를 두 경우로 나눈다. 이 점에서 $N$개의 ‘사건’을 지나면, $2^N$가지의 경우가 만들어지지 않을까 하는 걱정을 할 수 있다. 실제로는 그럴 필요가 없다. $a_i \le b_i$가 언제나 성립하기 때문에, $N$개의 사건을 지나면 만들어지는 경우는 최대 $N+1$가지이다. 이는 최악의 경우를 만들어 보려는 시도를 통해서도 직관을 얻을 수 있다.

사건 하나를 지나면 두 경우로 나뉠 것이다 (이를테면, 초기 상태가 3 이상이었으면 5 증가, 3 미만이었으면 2 감소). 마찬가지로, 사건 두 개를 지나면 세 경우로 나뉠 것이다 (이를테면, 초기 상태 3 이상이었으면 5 증가, 1이상 3 미만이었으면 1 증가, 1 미만이었으면 2 감소). 사건 세 개를 지나면 네 경우로 나뉠 것이다… 이를 기록할 수 있다. 즉, 구간 $[i..j]$에 대해, 구간의 사건들을 모두 겪었을 때의 가능한 상태 변화량 (초기 상태에 따라 몇 개의 분기가 나올 것이다)를 세그먼트 트리의 정점에 기록한다.

위와 같이 만들어진 세그먼트 트리의 정점을 하나씩 가리고 결과를 얻으면 사건이 하나씩 없을 때의 답을 알 수 있다. 즉, $i$번째 사건이 없었을 때의 결과는 두 구간 $[1..i-1]$과 $[i+1..N]$의 결과를 합치면 얻을 수 있다.

이 풀이의 시간 복잡도는 각 질의당 $O(\log N)$이 아닌데, 왜냐하면 세그먼트 트리의 정점이 차지하는 메모리의 크기가 구간의 길이에 비례하기 때문이다. 세그먼트 트리 연습문제에서는 세그먼트 트리의 정점이 int, 그리고 Addition Robot 문제에서도 int[2][2]이었기 때문에 크기가 항상 일정했다. 하지만 이 문제에서는 정점의 크기가 구간의 길이 $L$에 대해 $O(L)$이며, 길이 $L$인 두 정점을 합치는 데에 드는 시간도 $O(L)$이기 때문에, 정점 하나를 업데이트하는 데에 걸리는 시간은 최대 $O(N)$이다. 대신, 결과를 얻는 질의는 이분 탐색을 이용하면 구간의 길이 $L$에 대해 $O(\log L)$에 할 수 있기 때문에, 탐색할 정점의 개수 $O(\log N)$을 곱해 $O(\log^2 N)$의 시간에 질의를 해결할 수 있다. 이 문제에서는 업데이트 질의가 없고 결과를 얻는 질의만 $N$번 수행하면 되므로 총 시간 복잡도는 $O(N \log^2 N)$이 된다.

추가적으로, 이 문제와 유사한 Cut Inequality Down 문제도 풀어보면 도움이 될 것이다.

정리

세그먼트 트리는 node와 연산 *이 다음 성질을 만족할 때 사용할 수 있다.

연산이 결합 법칙을 만족한다. 즉, $(AB)C = A(BC)$.
연산에 항등원이 존재한다. 즉, 다음 성질을 만족하는 $I$가 존재한다: 모든 $A$에 대해 $AI = IA = A$. 하지만, 항등원이 없을 때에도 (억지로 약속하고 쓰는 한이 있더라도) 언제나 항등원을 만들어 정의할 수 있다.

지연 갱신은 node에 적용할 업데이트 연산 $f(node, x)$ (이를테면, node에 $x$를 더하는 연산)가 추가적으로 다음 성질을 만족할 때 사용할 수 있다.

연산 $f$가 여러개 중첩되어도 $f$로 표현가능해야 한다. (이를테면, $x$를 더하고 $y$를 더하는 연산은 $x+y$ 를 더하는 연산으로 표현된다.)
물론, 구간 $[i..j]$에 한꺼번에 적용할 수 있는 연산이어야 한다. (이를테면, 길이 5인 구간에 한꺼번에 $x$를 더하면 구간의 모든 수들의 합은 $5x$ 증가할 것이다.)

수학적 엄밀함을 위해서는 두 링크 세그먼트 트리, 지연 갱신을 참조하라.

세그먼트 트리는 아주 다양한 방향으로 응용되지만 이 글에서는 다양한 형태의 객체를 세그먼트 트리 안에 넣는 방법의 응용만을 살폈다. 이 외에도 평면 스위핑 등의 기법을 사용하여 문제의 형태가 세그먼트 트리의 쿼리를 이용하는 형태의 변형이 이루어지는 등, 아주 많은 방향으로 활용되는 자료구조이기 때문에, 많은 문제를 풀어보며 활용 사례를 알아나가는 것은 큰 도움이 된다.

세그먼트 트리의 응용

세그먼트 트리

세그먼트 트리의 응용

문제 1: Addition Robot

문제 2: 없던 일처럼

정리

이동하기 4 (BOJ 18796)

Stern-Brocot Tree를 활용한 수론적 함수의 합 계산