[자료구조] 세그먼트 트리(Segment Tree)

세그먼트 트리란?

배열에서 구간에 대한 정보를 효율적으로 관리(구간합, 최솟값, 최댓값 등)하기 위한 트리 구조의 자료구조이다.

 

 

특정 배열의 구간합을 구한다고 해보자

인덱스	0	1	2	3	4	5	6	7	8
데이터	3	4	5	6	7	8	9	12	11

위와 같은 배열이 있다고 해보자

인덱스 2~7까지의 구간합을 구한다고 하면

인덱스	0	1	2	3	4	5	6	7	8
데이터	3	4	5	6	7	8	9	12	11

인덱스 2 ~ 7까지 구간을 돌며 합을 구하면 될 것이다. 결과로는 47이 나올 것이고, 시간 복잡도는 O(7-2)의 시간이 들게 된다. N개의 데이터의 구간합을 구한다면 O(N)이 될 것이다. 만약 N이 많이 커지고, 구간합을 구하는 횟수가 많아지면 이러한 방식은 매우 비효율적이기 때문에 더욱 빠른 방법이 필요하다.

 

세그먼트 트리

이진 트리의 구조를 지니며 구간의 합을 저장할 수 있다. O(logN)의 시간으로 구간합 등을 구할  수 있다.

위 배열의 구간 합을 세그먼트 트리로 저장한다면 이러한 형식일 것이다.

이러한 형태로 세그먼트는 저장되어 있다. 인덱스 2~7까지의 구간합은 다음과 같이 구할 수 있다. 시간 복잡도는 O(logN)이다.

 

빨갛게 칠해져 있는 노드들의 합이 구간의 합이 될 것이다.

 

 

구현

구간합을 저장하는 세그먼트 트리의 코드는 다음과 같다.

//init

int arr[1000];
int seg[1000 * 4];

int init(int node, int start, int end){
	if(start == end){
        return seg[node] = arr[start];
    }
	int mid = s + (e - s) / 2;
    return seg[x] = init(node * 2, start, mid) + init(node * 2 + 1, mid + 1, end);
}

 

 

다음은 세그먼트 트리를 조회하는 함수이다. left~right까지 범위의 구간합을 O(logN)으로 구할 수 있다.

int query(int node, int start, int end, int left, int right){
	//구간을 벗어나는 범위는 필요없다.
	if(right < start || left > end) return 0;
    if(left <= start || end <= right) return seg[node];
    int mid = start + (end - start) / 2;
    return query(node * 2, start, mid, left, right) + query(node * 2 + 1, mid + 1, end, left, right);
}

 

 

다음은 idx번째 요소를 diff로 바꾸는 작업을 하는 함수이다.

int update(int node, int start, int end, int idx, int diff){
	//범위를 벗어나는 경우
    if(idx < s || idx > e) return 0;
    if(start == end){
    	seg[x] = diff;
        return;
    }
    int mid = start + (end - start) / 2;
    return seg[x] = update(node * 2, start, mid, idx, diff) + update(node*2+1,mid + 1, end, idx, diff);
}

 

위 세 함수로 기본적인 세그먼트 트리를 구성하고, 필요한 정보들을 저장해 사용할 수 있다.

 

 

 

참고

기본적인 세그먼트 트리는 update함수를 통해 한 번에 한 리프노드의 값만 변경이 가능하다. 하지만 구간 전체를 업데이트를 하고 싶은 경우에는 어떨까?

예를 들어 left ~ right까지 모든 노드에 a를 더하는 작업을 한다고 해보자.
이때는 느리게 갱신되는 세그먼트트리라는 응용이 필요하다. 세그먼트 트리에서는 구간합, 업데이트 등 노드에 접근할 때, 모든 노드에 접근하지 않는다. 그러므로 변화량인 a를 다른 한 곳에 저장해두고, 해당 노드를 접근할때  세그먼트 트리를 갱신하는 방식으로 하면 해결할 수 있을 것이다. 이를 느리게 갱신되는 세그먼트 트리라고 한다.

 

느리게 갱신되는 세그먼트 트리에서의 업데이트

int seg[100000 * 4];
int lazy[100000 * 4];

void update_lazy(int node, int start, int end){
	if(lazy[node] == 0) return;
    seg[node] += lazy(node) * (end - start + 1);
    
    //리프노드가 아닌 경우
    if(start != end){
    	lazy[node * 2] += lazy[node];
        lazy[node * 2 + 1] += lazy[node];
    }
    lazy[node] = 0;
}


void update(int node, int start, int end, int left, int right, int diff){
	update_lazy(node ,start, end);
    if(right < start || end < left) return;
    if(left <= start && end <= right){
    	lazy[node] = diff;
        update_lazy(node, star, end);
        return;
    }
    int mid = start + (end - start) / 2;
    update(node * 2, start, mid, left ,right, diff);
    update(node * 2 + 1, mid + 1, end, left ,right, diff);
    seg[node] = seg[node * 2] + seg[node * 2 + 1];
}



int query(int node, int start, int end, int left, int right){
	update_lazy(node, start, end);
    ....
}

// 세그먼트 트리에서의 업데이트는 모든 구간을 직접 순회하지 않는다. 업데이트가 모두 일어나지 않는다. 때문에 해당 아이디어는 노드에 접근할때 저장된 변화량을 활용해서 업데이트하는 것이다.

 

 

추천 문제

구간 합 구하기 https://www.acmicpc.net/problem/2042

최솟값과 최댓값 https://www.acmicpc.net/problem/2357

구간 곱 구하기 https://www.acmicpc.net/problem/11505

 

느리게 갱신되는 세그먼트 트리

구간 합 구하기2 https://www.acmicpc.net/problem/10999

수열과 쿼리 21 https://www.acmicpc.net/problem/16975