158. Segment Tree

알고리즘 문제를 풀 때 가장 자주 등장하는 자료 구조 중 하나인 세그먼트 트리(Segment Tree)에 대해 다룬다. 먼저 누적합을 다루면서 소개했던 구간 합 구하기 4 문제를 다시 보자.

수 $n$개가 주어졌을 때, $i$번째 수부터 $j$번째 수까지 합을 구하는 프로그램을 작성하시오.

 

이 문제에서 주어지는 쿼리는 아래와 같다.

• $i\ j$: $i$번째 수부터 $j$번째 수까지의 합을 출력한다.

수의 개수 $n$과 쿼리의 개수 $m$의 최댓값이 $10^5$이므로 누적합을 전부 구한 다음 각 쿼리를 $\Theta(1)$에 처리하면 $\Theta(n + m)$에 전체 문제를 해결할 수 있었다.

이제 쿼리를 하나 더 추가하자.

• $1\ x\ y$: $x$번째 수를 $y$로 바꾼다.
• $2\ x\ y$: $x$번째 수부터 $y$번째 수까지의 합을 출력한다.

원소의 값을 바꾸는 쿼리가 추가되었다. 이 쿼리들을 가장 간단하게 처리한다면 아래의 두 가지 방법 정도가 될 것이다.

• 각 원소의 값만 저장하고, $2$번 쿼리를 처리할 때 합을 직접 계산한다.
• 누적합을 저장하고, $1$번 쿼리를 처리할 때 누적합을 갱신한다.

첫 번째 방법에서 합을 계산하는 시간이나 두 번째 방법에서 누적합을 갱신하는 시간은 모두 $O(n)$이고, 이 쿼리가 최대 $m$개 들어올 수 있다. 따라서 두 방법의 시간복잡도는 모두 $O(nm)$이다. $n$과 $m$의 최댓값이 $10^5$이므로 문제를 풀 때 쓰기에는 적절하지 않다.

---

이제 합 세그먼트 트리를 사용하여 시간복잡도를 줄이는 방법을 알아보자. 먼저 합 세그먼트 트리의 구조는 다음과 같다.

$n = 8$이고 수들의 초기값이 리프 노드에 적힌 것과 같다고 하자. 세그먼트 트리는 루트가 있는 완전 이진 트리의 형태를 지니고 있으며 리프 노드에는 각 수의 값이, 나머지 노드에는 자식 노드에 적힌 수의 합이 대응된다.

 

쿼리는 다음과 같은 방법으로 처리한다. $x$번째 수를 $y$로 바꿀 경우, 리프 노드 하나와 조상 노드들의 값이 변하므로 $\Theta(\lg n)$개 노드의 값을 업데이트한다. 예를 들어 $3$번째 수를 $2$로 바꿀 경우 다음과 같이 $4$개 노드의 값을 업데이트한다.

$x$번째 수부터 $y$번째 수까지의 합을 구할 때는 루트 노드에서 탐색을 시작하여 해당 노드에 대응하는 구간이 $x$번째부터 $y$번째까지에 완전히 포함되면 그 노드의 값을 더하고, 일부 구간만 $x$번째부터 $y$번째까지에 포함되면 양쪽 자식 노드를 탐색하고, 하나도 겹치지 않으면 더 탐색하지 않는다. $2$번째 수부터 $5$번째 수까지의 합을 구하는 경우를 예시로 들어 보자. 파란색은 현재 탐색하는 노드, 초록색은 이후에 탐색해야 하는 노드, 빨간색은 값을 더한 노드이다.

루트 노드에서 탐색을 시작한다. 이 노드에 대응하는 구간은 $[1, 8]$이다. 수들의 합을 구하는 구간은 $[2, 5]$인데, $[1, 8]$의 일부만 $[2, 5]$에 포함되므로 양쪽 자식 노드를 탐색해야 한다.

자식 노드의 탐색 순서는 상관이 없다. 이 글에서는 왼쪽 노드부터 탐색하였다. 이번에 탐색하는 노드에 대응하는 구간은 $[1, 4]$이다. 다시 양쪽 자식 노드를 탐색한다.

이 노드에 대응하는 구간은 $[1, 2]$이다. 다시 양쪽 자식 노드를 탐색한다.

이 노드에 대응하는 구간은 $[1, 1]$이다. 이 구간은 $[2, 5]$와 겹치지 않으므로 더이상 탐색하지 않는다. 오른쪽 노드로 넘어간다.

이 노드에 대응하는 구간은 $[2, 2]$이다. 이 구간은 $[2, 2]$에 완전히 포함되므로 합에 $4$를 더하고, 부모 노드로 돌아간다.

다음에 탐색할 노드에 대응하는 구간은 $[3, 4]$이다. 이 구간은 $[2, 5]$에 완전히 포함되므로 합에 $9$를 더하고 부모 노드로 돌아간다.

다음에 탐색할 노드에 대응하는 구간은 $[5, 8]$이다. 양쪽 자식 노드를 탐색한다.

이 노드에 대응하는 구간은 $[5, 6]$이다. 양쪽 자식 노드를 탐색한다.

이 노드에 대응하는 구간은 $[5, 5]$이다. 이 구간은 $[2, 5]$에 완전히 포함되므로 합에 $1$을 더하고 오른쪽 노드로 넘어간다.

이 노드에 대응하는 구간은 $[6, 6]$이다. 이 구간은 $[2, 5]$와 겹치지 않으므로 부모 노드로 돌아간다.

다음에 탐색할 노드에 대응하는 구간은 $[7, 8]$이다. 이 구간은 $[2, 5]$와 겹치지 않으므로 부모 노드로 돌아간다.

더이상 탐색할 노드가 없으므로 탐색을 종료한다. $2$번째 수부터 $5$번째 수까지의 합은 $14$이다. 이 방법으로 합을 구할 경우 각 층에서 양쪽 자식을 탐색해야 하는 노드는 최대 $2$개이다. 즉 각 층에서 탐색하는 노드는 최대 $4$개이다. 층의 수가 $\lceil \lg n \rceil + 2$이므로 탐색하는 노드의 수는 $O(\lg n)$이다.

 

두 쿼리를 모두 $O(\lg n)$ 시간에 처리할 수 있으므로 $m$개의 쿼리를 처리하는 시간은 $O(m \lg n)$이 된다. 입력을 포함하면 전체 수행 시간은 $O(n + m \lg n)$이 되고, $n$과 $m$이 $10^5$이어도 충분히 빠르게 작동한다.

---

위에서 합 세그먼트 트리를 예시로 들었는데, 이것과 최댓값 세그먼트 트리, 최솟값 세그먼트 트리가 가장 많이 사용된다. 트리를 구성하는 방식은 합 세그먼트 트리와 비슷하다. 또한 AND, OR, XOR 등 비트 연산의 결과를 처리하는 세그먼트 트리도 가끔씩 등장한다. 그 외에 다른 응용도 가능하며, 몇 가지는 별도로 살펴볼 예정이다.

 

또한 원소의 개수가 $2$의 거듭제곱이 아닐 수도 있는데, 이때는 사용하지 않는 리프 노드가 생기게 된다. 이런 경우에도 나머지 노드의 값을 동일한 방식으로 저장하면 된다.

 

이제 구현 방법을 살펴보자. 세그먼트 트리를 구현할 때는 일반적으로 루트를 $1$번 노드, $x$번 노드의 자식을 $2x$번 노드와 $(2x+1)$번 노드로 두어서 연속된 공간에 노드의 정보를 저장하는 경우가 많다. 즉 노드에 다음과 같은 순서로 번호를 매기고 이 순서대로 연속된 공간에 저장하는 것이다.

이번에도 합 세그먼트 트리를 예시로 들 것이다. 먼저 각 노드에 대응하는 구간을 설정한다.

int p = 1 << 17, l[1 << 18], r[1 << 18];

for(int i = 0; i < p; i++)l[p + i] = r[p + i] = i;
for(int i = p; --i;)
{
    l[i] = l[i << 1];
    r[i] = r[i << 1 | 1];
} // i번 노드에 대응하는 구간은 [l[i], r[i]]가 된다.

직관적인 구조를 위해 첫 번째 리프 노드가 $0$번 원소를 가리키게 바뀌었다. 다음으로 원소들의 초기값을 설정한다.

int n, a[1 << 18];

cin >> n;
for(int i = 1; i <= n; i++)cin >> a[p + i];
for(int i = p; --i;)a[i] = a[i << 1] + a[i << 1 | 1];

다음으로 쿼리를 처리하는 함수를 구현한다. 먼저 값을 업데이트하는 쿼리는 다음과 같이 처리할 수 있다. 아래 코드의 경우 반복문의 가독성이 약간 떨어질 수 있는데, 쉽게 풀어서 구현해도 무방하다.

void REP(int i, int v) // i번째 값을 v로 바꾼다.
{
    for(a[i += p] = v; i >>= 1;)a[i] = a[i << 1] + a[i << 1 | 1];
}

구간 합을 구하는 쿼리는 크게 두 가지 방법으로 처리하는데, 하나는 재귀함수를 사용하고 하나는 재귀함수를 사용하지 않는다. 재귀함수를 사용하지 않는 방식이 속도가 더 빠르지만, 세그먼트 트리를 처음 접하는 사람에게 재귀함수를 사용하지 않는 방식부터 설명하는 것은 그다지 좋은 생각이 아니기 때문에 재귀함수를 사용하지 않는 방식은 이후에 따로 다루고 여기서는 재귀함수를 사용하는 방식을 소개하기로 한다.

int getSum(int x, int y, int d) // x번째 원소부터 y번째 원소까지의 합을 구해야 하고, 현재 탐색하는 노드가 d번인 경우
{
    if(x <= l[d] && r[d] <= y)return a[d];
    if(x > r[d] || y < l[d])return 0;
    return getSum(x, y, d << 1) + getSum(x, y, d << 1 | 1);
}

리프 노드를 탐색할 경우, 그 노드에 대응하는 구간의 길이는 $1$이므로 합을 구해야 하는 구간에 완전히 포함되거나 하나도 포함되지 않는 상황만 발생하고 두 자식 노드를 탐색하는 상황은 발생하지 않는다. 따라서 이 재귀호출은 정상적으로 실행된다.

 

[연습문제]

 

BOJ 2042. 구간 합 구하기 (Gold I)

합 세그먼트 트리를 이용하여 쿼리를 처리한다.

 

BOJ 2357. 최솟값과 최댓값 (Gold I)

최솟값 세그먼트 트리와 최댓값 세그먼트 트리를 이용하여 쿼리를 처리한다.

 

BOJ 11505. 구간 곱 구하기 (Gold I)

곱 세그먼트 트리를 이용하여 쿼리를 처리한다. 곱 세그먼트 트리는 일반적으로 잘 등장하지 않는 편이다.

 

→ solved.ac tag: segtree