세그먼트 트리 (Segment Tree) - C/C++로 설명

세그먼트 트리(Segment Tree)는 효율적으로 구간합이나 최소값, 최대값을 구하는 데 매우 유용한 자료구조이다. 특히, 데이터가 자주 업데이트되면서도 특정 구간의 값을 빠르게 조회해야 하는 상황에서 효과적이다. 다음은 세그먼트 트리의 기본 개념과 사용법을 예제를 통해 정리한 내용이다.

세그먼트 트리 - 쓰임새

세그먼트 트리는 연속적인 데이터가 있을 때, 특정 범위의 합/ 최소, 최대값 등을 구할 때 유용하게 활용될 수 있다.

특정 구간의 합을 구하는 방법에 대해 살펴보자면, 일반적으로 다음과 같은 방법이 떠오른다..

 

  1. $arr[l] + arr[l+1] + ... + arr[r-1] + arr[r]$을 일일히 더해 구하는 방법

  2. $i$번째까지의 합을 저장하는 배열을 하나 더 만들어서, 조금 더 손쉽게 계산하는 방법

 

for(int i = 1; i <= n; i++) {
	cin >> arr[i];
	sum[i] = sum[i-1] + arr[i];
}

 

2번 방법의 경우, 꽤 효율적이라 할 수 있겠지만 `arr` 배열 안의 값에 변동이 생길 때마다 `sum` 배열을 갱신해줘야 한다. 일례로 맨 앞의 원소가 바뀌는 경우에는, `sum` 배열의 모든 값을 바꿔야 한다. 이렇게 값의 변동이 이루어지는 상태에서 연속적으로 구한 합을 구해야 하는 문제가 주어진다면, 2번 방법도 그리 좋은 선택지가 되지 못한다. 

 

세그먼트 트리는 이와 같은 상황에서 효과적이다.

 

만약 $i$번째 수를 바꾸고, 구간합을 구하는 연산을 $M$번씩 하게 된다면 위 1, 2번 방식으로는 $O(MN)$의 시간 복잡도가 소요되는 반면,

( 구간합 구하는데 $O(N)$, 원소를 바꾸는데 $O(1)$ → $O(MN + N) = O(MN)$ ) 

 

세그먼트 트리는 이 연산을 $O(MlogN)$으로 처리할 수 있다.

( 구간합 구하는데 $O(logN)$, 원소를 바꾸는데 $O(logN)$ →  $O(MlogN+MlogN)=O(MlogN)$ )

 

세그먼트 트리 - 개형

세그먼트 트리의 개형을 알아보자.

 

다음과 같이 트리의 leaf 노드(가장 밑바닥)에는 원래 배열의 값이 순서대로 오고, 각각의 노드의 부모 노드에는 자식 노드들의 합이 오게 된다. (설명의 편의를 위해 완전 이진 트리로 나타내었다.)

 

배열의 세그먼트 트리 표현

 

노드의 인덱스는

루트 노드가 1,

왼쪽 자식 노드는 $(부모 노드의 인덱스)\times 2$

오른쪽 자식 노드는 $(부모 노드의 인덱스)\times 2 + 1$로 설정한다. 이와 같이 설정했을 때 재귀적으로 트리를 나타내고 계산하기 수월해진다.

아래 그림으로 다시 트리를 살펴보자. 보라색이 i - j번째의 합을 나타내는 구간의 정보, 붉은색이 각 노드의 인덱스이다. 헷갈리지 말자.

 

세그먼트 트리의 인덱스

트리의 초기화

최종적으로 위와 같은 현태의 트리를 만들기 위해서, 다음과 같은 재귀함수 꼴의 초기화 함수를 사용한다. 물론 반복문 형태로도 짤 수 있으나 재귀함수로 나타내는 것이 더 직관적이다.

(`arr`: 원래 수가 저장되어 있는 배열, `segTree`: 세그먼트 트리를 나타내는 배열)

 

int init(int node, int st, int end) {
    // st   : init 함수가 관심두는 arr의 시작 인덱스
    // end  : init 함수가 관심두는 arr의 끝 인덱스
    // node : segTree의 노드
    // -> node번째 노드가 st ~ end의 합을 저장한다.
    if (st == end) return segTree[node] = arr[st];
    int mid = (st + end) / 2;

    //재귀로 반씩 나눠서 초기화
    return segTree[node] = init(node * 2, st, mid) + init(node * 2 + 1, mid + 1, end);
}

 

이때 `segTree` 배열의 길이가 궁금할 수 있다. 세그먼트 트리의 노드의 개수는 다음과 같이 구할 수 있다.

  1. 트리의 높이 `h`는 $log2( leaf 노드의 수(원래 arr 배열의 원소 수)$ )를 올림하여 구할 수 있다.

  2. 등비수열의 합 공식에 의해, $2^{(h+1)}$이 된다.

 

하지만 이러한 계산 과정이 복잡하다 느낀다면, 조금 메모리 공간을 잡아먹더라도 $4 × (arr 배열의 원소 수)$로 처리할 수 있다.

 

원소 업데이트(갱신)

`init` 함수에 의해 만들어진 트리에서 특정 원소를 바꾸고자 할 땐 다음과 같이 `update` 함수를 짤 수 있다.

 

void update(int n, int st, int end, int t, int diff) {
    // st   : 시작 인덱스
    // end  : 끝 인덱스
    // idx  : 수정할 원소의 인덱스
    // diff  : 수정할 값

    // 범위 안에 있을 경우
    if (st <= t && t <= end) segTree[n] += diff;

    // 범위 밖에 있을 경우
    else return;
    if (st == end) return;
    int mid = (st + end) / 2;
    update(n * 2, st, mid, t, diff);
    update(n * 2 + 1, mid + 1, end, t, diff);
}

 

`init` 함수와 마찬가지의 재귀함수 형태이다.

 

구간 합

마지막으로, 구간 합을 구하는 `sum` 함수를 작성할 수 있다. 

 

int sum(int l, int r, int node, int st, int end) {
    // st   : 시작 인덱스
    // end  : 끝 인덱스
    // l~r  : 구하고자 하는 구간 합의 범위

    // [l, r]이 [st, end]를 완전히 포함하는 경우
    if (l <= st && end <= r) return segTree[node];

    // [l, r]와 [st, end]가 겹치지 않는 경우
    if (r < st || end < l) return 0;

    // 나머지 경우 (일부분 겹칠때)
    int m = (st + end) / 2;
    return sum(l, r, node * 2, st, m) + sum(l, r, node * 2 + 1, m + 1, end);
}

 

Summary

init(1, 1, N);

와 같이 작성하여 트리 초기화를 할 수 있고, 

 

update(1, 1, N, t, diff);

로 `t`번째 원소 값을 `diff`만큼 바꿀 수 있다.

 

sum(l, r, 1, 1, N)

는 `l~r`번째 원소의 구간합을 구하는 호출이 될 것이다.

 

 

세그먼트 트리의 구현 방법에 대해 개략적으로 알아보았다면 다음 백준 예제 풀이를 통해 연습해보자: 

[C++] 백준 2042 구간 합 구하기