세그먼트 트리

구간 합을 구할 때 사용하는 유용한 자료구조인 세그먼트 트리다. 구간 합은 데이터가 나열되어 있을 때, 특정 범위에서의 데이터 합을 말한다. 누적합을 구했다는 가정 하에 구간 합은 O(1)의 시간복잡도로 구할 수 있지만, 변경된 사항이 생긴다면 다시 O(N)의 시간복잡도로 누적합을 수정해야 한다. 하지만 세그먼트 트리를 이용하면 수정했을 때도 트리의 높이인 O(lgN)만으로 누적합을 구할 수 있다

 

1. 세그먼트 트리(누적합) 만들기

리프 노드를 배열 원소를 하고 루트로 올라갈수록 재귀적으로 누적합이 계산된다. 완전 이진트리로 구성하기 때문에 자식 인덱스는 2*부모 + 1, 2*부모 + 2로 구성된다

 

https://wkdtjsgur100.github.io/segment-tree/

 

https://wkdtjsgur100.github.io/segment-tree/

 

long long init(int idx, int start, int end) {
	if (start == end) {
		return tree[idx] = arr[start];
	}
	else {
		return tree[idx] = init(2 * idx + 1, start, (start + end) / 2) +
			init(2 * idx + 2, (start + end) / 2 + 1, end);
	}
}

 

2. 구간합 구하기

일정 구간이 주어졌을 때, 기본으로 만들어진 세그먼트 트리와 비교했을 때 3가지로 나뉠 수 있다

 

1. 범위가 하나도 겹치지 않을 때

2. 일정 구간이 세그먼트 트리 구간을 포함할 때

3. 그 외

 

long long sum(int idx, int start, int end, int left, int right) {
	if (left > end || right < start) {
		return 0;
	}
	else if (left <= start && end <= right) {
		return tree[idx];
	}
	else {
		int mid = (start + end) / 2;
		return sum(2 * idx + 1, start, mid, left, right) + 
			sum(2 * idx + 2, mid + 1, end, left, right);
	}
}

 

3. 값 변경했을 때, 구간합 변경하기

변경하려는 인덱스가 포함된 구간합은 모두 차이만큼 더해준다 

 

void change(int idx, ll diff, int start, int end, int changeidx) {
	if (start > changeidx || end < changeidx) return;

	tree[idx] += diff;
	if (start < end) {
		int mid = (start + end) / 2;
		change(2 * idx + 1, diff, start, mid, changeidx);
		change(2 * idx + 2, diff, mid + 1, end, changeidx);
	}
}