[Algorithm] 세그먼트 트리(Segment Tree)

구간합, 구간에 따른 경우의 수를 구하는 문제에서 많이 쓰이는 알고리즘인 세그먼트 트리에 대해 알아보자.

 

* 세그먼트 트리의 등장 배경 및 필요성

S[0] = A[0];
for (int i=1; i<n; i++) {
    S[i] = S[i-1] + A[i];
}

 

구간합을 단순 for 문을 이용해서 구하는 코드이다.

0~n 까지의 합을 구할 때 O(n) 의 시간복잡도가 발생한다.

0~n 뿐 아니라 2~5, 100~200 등 m 개의 구간합을 구한다면 O(nm) 의 시간복잡도가 발생한다.

 

여기서 만약 A[x] 의 값을 변경해야 하면 어떻게 될까.

A[x] 값이 포함된 모든 S 배열 값을 변경해서 다시 구간합을 구해주어야 한다.

0번 값이라고 하면 또 O(nm) 의 시간복잡도가 발생한다.

 

큰 수의 n, m 이 주어지면 시간초과로 풀 수 없다. 다른 방법을 찾아야 한다.

---

* 세그먼트 트리란 ?

: 구간을 저장하기 위한 트리.

 

- 세그먼트 트리의 리프노드: 배열의 그 수 자체

- 세그먼트 트리의 다른 노드: 왼쪽 자식과 오른쪽 자식의 합

- 한 노드의 번호가 x일 때 왼쪽 자식의 번호는 2*x, 오른쪽 자식의 번호는 2*x + 1 이 된다.

 

n = 10 인 배열을 세그먼트 트리로 표현해보자.

출처: https://www.acmicpc.net/blog/view/9

 

리프 노드에 각 배열의 index 가 들어 있고, 점점 루트 노드로 가며 0~1의 합, 0~1과 2의 합, 3과 4의 합, 0~2와 3~4의 합, 0~4와 5~9의 합이 담긴다. 시간복잡도 O(logN).

 

왼쪽, 오른쪽 트리의 나눔 기준은 배열 크기의 반이다.(위에 적은 세그먼트 트리의 특징으로 한 노드의 번호가 x일 때 왼쪽 자식의 번호는 2*x, 오른쪽 자식의 번호는 2*x + 1 이기 때문.)

 

즉, 0~9의 구간합을 구해야 한다면 루트노드만 접근하면 된다.

5~8의 구간 합을 구해야 하면 아래와 같이 2개의 노드만 접근하면 된다.

출처: https://www.acmicpc.net/blog/view/9

 

3~9의 구간합을 구해야 하는 경우도 아래와 같이 2개의 노드만 접근하면 된다.

출처: https://www.acmicpc.net/blog/view/9

 

배열의 어떤 값을 변경해야 할 때도 해당 값을 포함하는 값과 구간 노드만 변경해주면 된다.

출처: https://www.acmicpc.net/blog/view/9

세그먼트 트리가 아니었다면 모든 구간합에 접근하여 값을 변경해줘야 했을 것이다. 세그먼트 트리의 특성을 이용하면 4개의 구간합에만 업데이트하면 된다. 시간복잡도 O(logN).

 

 

* 세그먼트 트리 만드는 법(C++)

1) 세그먼트 트리 생성

// a: 배열 a
// tree: 세그먼트 트리
// node: 세그먼트 트리 노드 번호
// node가 담당하는 합의 범위가 start ~ end
long long init(vector<long long> &a, vector<long long> &tree, int node, int start, int end){
    if(start == end){
    	return tree[node] = a[start];
    } else{
    	return tree[node] = init(a, tree, node*2, start, (start+end)/2) + init(a, tree, node*2+1, (start+end)/2+1, end);
    }
}

start == end 인 경우는 리프노드를 뜻하므로 리프노드는 배열 자체의 값이 포함되어야 하므로 tree[node] 에 a[start] 값을 넣어준다.

리프노드가 아닌 경우엔 tree[node] 에 왼쪽 자식 노드의 값과 오른쪽 자식 노드의 값을 더해서 넣어준다.

 

2) 세그먼트 트리 구간합 구하기

// node 가 담당하는 구간: start ~ end
// 구해야하는 합의 범위: left ~ right
long long sum(vector<long long> &tree, int node, int start, int end, int left, int right) {
     if(left > end || right < start) {
     	return 0;
     }
     if(left <= start && end <= right) {
     	return tree[node];
     }
     return sum(tree, node*2, start, (start+end)/2, left, right) + sum(tree, node*2+1, start+end)/2+1, end, left, right);
}

 

구해야하는 구간합 [left, right] 의 범위가 [start, end] 에 하나라도 없다면 0을 return 한다.

구해야하는 구간합 [left, right] 의 범위가 [start, end] 안에 완전히 있다면 해당 node 값을 바로 return 한다.

그 경우가 아니라면 왼쪽, 오른쪽 노드 탐색을 통해 값을 구한다.

 

3) 세그먼트 트리 노드 값 변경하기

// tree: 세그먼트 트리
// node: 세그먼트 트리의 현재 노드 번호
// node 가 담당하는 합의 범위 start ~ end
// index: 수정해야할 노드의 index
// diff: 기존 값과 변경할 값의 차이
void update(vector<long long> &tree, int node, int start, int end, int index, long long diff) {
    if (index < start || index > end) return;
    tree[node] = tree[node] + diff;
    if (start != end) {
        update(tree,node*2, start, (start+end)/2, index, diff);
        update(tree,node*2+1, (start+end)/2+1, end, index, diff);
    }
}

index 가 start 와 end 사이의 값이 아니라면 바로 return 한다.

index 가 start 와 end 사이의 값이라면 diff 만큼을 더해서 값을 갱신한다.

leef 노드가 아니면 왼쪽, 오른쪽 자식 트리의 노드를 update 하러 간다.

---

* 출처 : https://www.acmicpc.net/blog/view/9

---

* 세그먼트 트리로 풀 수 있는 문제

2042번 구간합(https://www.acmicpc.net/problem/2042)

#include <cstdio>
#include <cmath>
#include <vector>
using namespace std;

long long init(vector<long long> &a, vector<long long> &tree, int node, int start, int end) {
    if (start == end) {
        return tree[node] = a[start];
    } else {
        return tree[node] = init(a, tree, node*2, start, (start+end)/2) + init(a, tree, node*2+1, (start+end)/2+1, end);
    }
}
void update(vector<long long> &tree, int node, int start, int end, int index, long long diff) {
    if (index < start || index > end) return;
    tree[node] = tree[node] + diff;
    if (start != end) {
        update(tree,node*2, start, (start+end)/2, index, diff);
        update(tree,node*2+1, (start+end)/2+1, end, index, diff);
    }
}
long long sum(vector<long long> &tree, int node, int start, int end, int left, int right) {
    if (left > end || right < start) {
        return 0;
    }
    if (left <= start && end <= right) {
        return tree[node];
    }
    return sum(tree, node*2, start, (start+end)/2, left, right) + sum(tree, node*2+1, (start+end)/2+1, end, left, right);
}
int main() {
    int n, m, k;
    scanf("%d %d %d",&n,&m,&k);
    vector<long long> a(n);
    int h = (int)ceil(log2(n));
    int tree_size = (1 << (h+1));
    vector<long long> tree(tree_size);
    m += k;
    for (int i=0; i<n; i++) {
        scanf("%lld",&a[i]);
    }
    init(a, tree, 1, 0, n-1);
    while (m--) {
        int t1,t2,t3;
        scanf("%d",&t1);
        if (t1 == 1) {
            int t2;
            long long t3;
            scanf("%d %lld",&t2,&t3);
            t2-=1;
            long long diff = t3-a[t2];
            a[t2] = t3;
            update(tree, 1, 0, n-1, t2, diff);
        } else if (t1 == 2) {
            int t2,t3;
            scanf("%d %d",&t2,&t3);
            printf("%lld\n",sum(tree, 1, 0, n-1, t2-1, t3-1));
        }
    }
    return 0;
}