[자료구조] 백준 2042번: 구간 합 구하기 (C++), 세그먼트 트리

https://www.acmicpc.net/problem/2042

2042번: 구간 합 구하기

일반적으로라면 구간 합을 구하는 영역이 3 ~ 7로 주어졌을 때, 3부터 7까지 for문을 써서 3~7 사이의 합을 구하는 것이 일반적이었다. 하지만, 이 문제에서 그런 방법을 사용한다면 시간 복잡도는 O(N * K)로 시간 초과를 유발한다.

 

오늘 공부한 세그먼트 트리는 시간 복잡도 O(N * K)에서 K를 logK로 줄여주는 신기한 현상을 보여준다.

 

공부한 블로그

(https://m.blog.naver.com/ndb796/221282210534)

(https://eun-jeong.tistory.com/18)

(21.12.13 구글링하다가 좋은 글 하나 더 발견 구간 합에 대하여 https://hongjw1938.tistory.com/20)

[세그먼트 트리 쓰는 이유]

특정 구간의 합, 최댓값, 최솟값 등 특정 구간에서의 어떤 값을 구할 때 빠르게 구할 수 있는 트리이다.

트리에 노드 별로 각 구간의 값들을 저장해두고서 어떠한 구간이 필요할 때, 보다 빠르게 가져올 수 있다.

[세그먼트 트리 정리]

0) 세그먼트 트리 크기

일반적으로는 배열의 크기 X4를 하여 선언하는데 조금의 메모리 낭비를 할 수는 있다만 큰 생각을 들이지 않는다면 저 방법이 제일 편하다.

더 디테일하게 구할 수 있는 방법은 있다.

N이 2의 제곱꼴인 경우에는 Full Binary Tree가 된다. 그러므로, 총 필요한 노드의 개수는 2^N - 1개이다.

N이 2의 제곱꼴이 아닌 경우에는 트리의 높이가 h = [logN]이 되고, 필요한 노드의 개수는 2^(h + 1) - 1이 된다.

int h = (int)ceil(log2(n)); // 트리의 높이
int tree_size = (1 << (h + 1)); // 2^(h + 1)

 

1) 세그먼트 트리 시간 복잡도

(https://eun-jeong.tistory.com/18)

초기화 과정: 각 노드마다 O(1)이므로 O(N)

쿼리 과정: O(logN)

업데이트 (노드 1개): 노드 1개를 포함하는 구간은 트리에 logN개 있기 때문에 O(logN)

업데이트 (구간):  노드 하나를 업데이트 할 때 O(logN)이기 때문에 최대 O(NlogN)

>> 업데이트 (구간)을 O(logN)으로 줄일 수 있는 Lazy Propagation (O(logN)) 존재

 

2) 구간 합 트리 생성

세그먼트 트리의 첫 노드 인덱스는 배열처럼 0이 아닌 1부터 시작해야 한다.

node * 2 혹은 node * 2 + 1 연산을 통해 자식 노드로 가게 해야 하는데 루트 노드가 0이라면 곤란해진다.

 

long long init(int start, int end, int node) {
	if (start == end) { // leaf 노드
		return tree[node] = arr[start];
	}
	int mid = (start + end) / 2;
	return tree[node] = init(start, mid, node * 2) + init(mid + 1, end, node * 2 + 1);
    // 부모 노드 = 자식 노드 + 자식 노드
}

 

3) 구간 합 구하기 (Sum), 배열의 값이 바뀔 때 (Update)

구간 합을 생성하는 세그먼트 트리가 이해가 됐다면 다음 구간 합을 구하는 함수와 배열의 값이 바뀌었을 때, 트리에는 어떻게 업데이트하는지를 이해할 수 있을 것이다.

#include <iostream>
#define MAX 1000000 + 1
using namespace std;

long long tree[4 * MAX];
long long arr[MAX];

int n, m, k;

//== Segment Tree INIT == (각 노드에 구간 합)
long long init(int start, int end, int node) {
	if (start == end) {
		return tree[node] = arr[start];
	}
	int mid = (start + end) / 2;
	return tree[node] = init(start, mid, node * 2) + init(mid + 1, end, node * 2 + 1);
}

long long sum(int start, int end, int node, int left, int right) {
	// 범위 밖에 있는 경우
	if (left > end || right < start) return 0;
	// 범위 안에 있는 경우
	if (left <= start && end <= right) return tree[node];
	// 그렇지 않다면 두 부분으로 나누어 합을 구하기
	int mid = (start + end) / 2;
	return sum(start, mid, node * 2, left, right) 
		+ sum(mid + 1, end, node * 2 + 1, left, right);
}

void update(int start, int end, int node, int index, long long dif) {
	if (index < start || index > end) return;
	tree[node] += dif;
	if (start == end) return;
	int mid = (start + end) / 2;
	update(start, mid, node * 2, index, dif);
	update(mid + 1, end, node * 2 + 1, index, dif);
}

int main() {
	cin >> n >> m >> k;
	for (int i = 1; i <= n; i++) {
		cin >> arr[i];
	}

	init(1, n, 1);
	for (int i = 0; i < m + k; i++) {
		int a, b;
		cin >> a >> b;
		long long c;
		cin >> c;
		if (a == 1) {
			long long dif = c - arr[b];
			arr[b] = c; // 이 부분은 없어도 된다. 어차피 세그먼트 트리 위주로 돌아가는 문제라서
			// 배열은 따로 업데이트 안 해줘도 됌.
			update(1, n, 1, b, dif);
		}
		else {
			cout << sum(1, n, 1, b, c) << '\n';
		}
	}

	return 0;
}