segment tree

(이 내용은 제가 이해한 걸 정리해둔 거라 틀린 내용이 있을 수도 있습니다.)

 

segment tree : 특정 구간에서 값이 변할 때 구간에 대한 쿼리를 처리하기 위한 자료구조(최대, 최소, 구간합, ... 등)

 

예를 들어 n개의 수열 a1, a2, ... an 이 있다고 했을 때 [l, r]구간의 합을 naive하게 구하면 O(r-l+1)의 시간복잡도를 가진다. d1=a1, dn=d(n-1)+an 이 점화식을 활용하여 구간합을 미리 구해두면 O(1)가 되지만 구간의 값이 바뀌게 된다고 하면 O(r-l+1)의 시간복잡도를 가지게 될 것이다. 이를 O(logN)만에 구해줄 자료구조를 segment tree(or fenwick tree(=binary indexed tree))라고 한다.

 

segement tree의 원리수열(1차원배열)을 트리 형태로 구성해서 구간합을 저장한다.

수열 a1, a2, ..., a7 가 있고, 트리의 각 노드가 t1, t2, ..., t15라고 할 때 각 수열의 값들은 리프 노드에 순서대로 대응된다. 여기서는 t8=a1, t9=a2, ..., t14=a7이고, t15는 대응되는 원소가 없으므로 0의 값을 가진다. 부모 노드는 모든 자식 노드 값들의 합으로 표현된다. 예를 들어 t4=t8+t9=a1+a2, t3=t6+t7=(t12+t13)+(t14+t15)=a5+a6+a7이라 할 수 있다. 여기서 자료의 개수가 늘어난다면 리프 노드를 더 많이 가지는 트리를 생각하면 된다. 자료의 개수가 n개라면 트리의 높이는 floor(log2(n))+1이고, 트리에서 필요한 노드의 수는 2^(floor(log2(n))+1)-1이다. segment tree는 비재귀형식과 재귀형식 두 가지로 구현할 수 있는데 이 그림은 비재귀로 구현했을 때라고 할 수 있다.

#include<iostream>
#include<cmath>
using namespace std;
using ll = long long;

ll tree[2100000];

int main() {
	ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	int n, m, k;
	cin >> n >> m >> k;
	int pos = (int)pow(2, ceil(log2(n))); //pos : 트리에 값을 저장할 때 첫번째 위치
	for (int i = 0; i < n; i++) {
		cin >> tree[pos + i];
		int index = pos + i;
		while (index) { // 부모노드(1)까지 값을 갱신해준다
			tree[index / 2] += tree[pos + i];
			index /= 2;
		}
	}
	
	ll a, b, c;
	for (int i = 0; i < m + k; i++) {
		cin >> a >> b >> c;
		if (a == 1) {
			int index = pos + b - 1;
			ll diff = c - tree[index]; // 차이를 구해서
			while (index) { // 부모노드(1)까지 값을 갱신해준다
				tree[index] += diff;
				index /= 2;
			}
		}
		else {
			int left = pos + b - 1;
			int right = pos + c - 1;
			ll sum = 0;
			while (left <= right) {
				if (left % 2 == 1) sum += tree[left++]; // left가 오른쪽 자식이면 값을 더함
				if (right % 2 == 0) sum += tree[right--]; // right가 왼쪽 자식이면 값을 더함
				left /= 2; // 부모노드로 올라간다
				right /= 2;
			}
			cout << sum << '\n'; // 구간합을 출력
		}
	}
	return 0;
}

 

이 그림은 7개의 값이 있을 때 segment tree를 재귀형식으로 구현한 것인데 비재귀와의 차이점은 값의 위치에 있다. 위의 그림과 비교해보면 7번 노드의 자식이 없다는 것인데 비재귀에선 a7을 t14에 저장하지만 재귀에서는 a7을 t7에 저장한다. 트리의 형태는 자료의 수가 2^(floor(log2(n))-1일 경우에는 재귀와 비재귀가 같아진다. 하는 일은 같지만 재귀형식으로 트리를 구현하면 lazy propagation(구간에 대한 값 갱신)이나 HLD(각 노드에 순서를 부여해서 트리를 일자로 편 뒤 트리의 구간에 대한 쿼리를 처리)등의 작업도 할 수 있다. 자신에게 편한 방법 하나만 사용하면 되지만 둘 다 알아두는 게 좋을 것 같다.

#include<iostream>
using namespace std;
using ll = long long;

ll ar[1000005];
ll tree[2100000];

ll init(int node, int left, int right) { // 처음에 입력받은 값들로 세그먼트트리를 생성
	if (left == right) return tree[node] = ar[left];
	return tree[node] = init(node * 2, left, (left + right) / 2) + init(node * 2 + 1, (left + right) / 2 + 1, right);
}

ll update(int node, int left, int right, int idx, ll value) { // idx위치의 값을 value로 변경
	if (idx < left || right < idx) return tree[node];
	if (left == right) return tree[node] = value;
	return tree[node] = update(node * 2, left, (left + right) / 2, idx, value) + update(node * 2 + 1, (left + right) / 2 + 1, right, idx, value);
}

ll query(int node, int nleft, int nright, int left, int right) { // nleft, nright : 구간을 탐색
	if (right < nleft || nright < left) return 0; // left, right : 원하는 구간
	if (left <= nleft && nright <= right) return tree[node];
	return query(node * 2, nleft, (nleft + nright) / 2, left, right) + query(node * 2 + 1, (nleft + nright) / 2 + 1, nright, left, right);
}

int main() {
	ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	int n, m, k;
	cin >> n >> m >> k;
	for (int i = 1; i <= n; i++) cin >> ar[i];
	init(1, 1, n);

	ll a, b, c;
	for (int i = 0; i < m + k; i++) {
		cin >> a >> b >> c;
		if (a == 1) update(1, 1, n, b, c);
		else cout << query(1, 1, n, b, c) << '\n';
	}
	return 0;
}

 

시간복잡도는 [l, r](1<=l<=r<=n)구간의 합을 구할 때 O(log(n))에 구할 수 있고(트리의 높이), ax(1<=x<=n)값이 바뀐다 해도 O(log(n))에 값 수정이 가능하다. 따라서 구간합, 값 변경 모두 O(log(n))이다.

 

segment트리와 비슷한 fenwick tree라는 자료구조도 있다.(=binary indexed tree = BIT) fenwick tree도 구간합을 구할 수 있어 segmenet tree와 비슷하지만 비트 연산을 사용한다는 것과, [l, r]의 구간합을 구할 때 [1, r]의 구간합 - [1, l-1]의 구간합으로 구한다는 차이점이 있다. fenwick tree도 값 수정, 구간합을 O(log(n))에 처리할 수 있다.