[백준] 2042번 구간 합 구하기

2042번: 구간 합 구하기

일반적인 방법으로 구간의 합을 구하면 O(N)이지만

이 문제에서는 시간 초과가 걸릴 가능성이 아주 높기 때문에

O(logN)인 세그먼트 트리를 이용하여 구해야 한다.

 

일차원 배열을 이용하여 이진 트리를 만드는 데

계산을 편하게 하기 위해서 루트를 인덱스 1로 잡고

왼쪽 노드는 부모 노드 * 2, 오른쪽 노드는 부모 노드 * 2 + 1

의 형태로 지정한다.

 

루트 노드를 반으로 나누어 구간 A, 구간 B가 되었다고 할 때

왼쪽 노드는 구간 A, 오른쪽 노드는 구간 B가 된다.

이 노드들이 다시 부모 노드가 되어 반으로 나누는 과정을

재귀 호출을 통해 반복하다가

구간의 길이가 1이 되면 그 구간의 값을 반환 한다.

 

그렇게 되면 거슬러 올라가서

각 노드는

구간 A의 합 + 구간 B의 합

이 계산된 상태가 된다.

 

1. 초기화(Init)

위의 방식으로 Init 함수를 만들어 초기 트리를 구성한다.

 

2. 값을 수정해야 하는 경우(Update)

Init 함수의 형태에서

수정해야하는 인덱스를 포함한 구간만 값을 변경시킨다는

조건을 추가한 Update 함수를 이용한다.

 

3. 출력하는 경우(Find)

1) 현재 노드의 범위가 목표 범위(출력해야 하는 범위)와 아예 겹치지 않을 때

→탐색 종료

 

2) 현재 노드의 범위가 목표 범위와 부분적으로 겹칠 때

→자식 노드를 탐색하러 감

 

3) 현재 노드의 범위가 목표 범위에 완전히 포함될 때

→결과에 현재 노드의 값 더함

 

이 세 가지 함수를 구현해주면 된다.

#include <iostream>
using namespace std;
int N, M, K, a, b;
long long arr[1000002], tree[4000008], res, c;
long long Init(int left, int right, int num) //왼쪽 구간, 오른쪽 구간, 현재 노드 넘버를 넘겨줌
{
	if (left == right) return tree[num] = arr[left]; //단말 노드는 자기 자신 반환
    
	int mid = (left + right) / 2;
	return tree[num] = Init(left, mid, num * 2) + Init(mid + 1, right, num * 2 + 1); //현재 노드는 왼쪽 노드와 오른쪽 노드의 값을 더한 값임
}
long long Update(int left, int right, int num, int idx, long long val) 
{
	if (left > idx || right < idx) return tree[num]; //바꿀 필요 없는 구간은 그대로 넘겨줌
    
	if (left == right) return tree[num] = val; //단말 노드는 바뀐 값 반환
    
	int mid = (left + right) / 2;
	return tree[num] = Update(left, mid, num * 2, idx, val) + Update(mid + 1, right, num * 2 + 1, idx, val);
}
void Find(int left, int right, int num, int target_left, int target_right) 
{
	if (left >= target_left && right <= target_right) res += tree[num]; //현재 범위가 목표 범위에 완전히 포함될 때 그 값을 가져옴
    
	else if (left > target_right || right < target_left) return; //현재 범위가 목표 범위와 아예 겹치지 않을 때는 멈춤
    
	else { //현재 범위가 목표 범위에 부분적으로 포함될 때는 자식 노드로 들어가서 찾음
		int mid = (left + right) / 2;
		Find(left, mid, num * 2, target_left, target_right);
		Find(mid + 1, right, num * 2 + 1, target_left, target_right);
	}
}
int main() {
	cin >> N >> M >> K;
	for (int i = 0; i < N; i++) cin >> arr[i];
	Init(0, N - 1, 1);
	for (int i = 0; i < M + K; i++) {
		cin >> a >> b >> c;
		if (a == 1) arr[b - 1] = c, Update(0, N - 1, 1, b - 1, c);
		else {
			Find(0, N - 1, 1, b - 1, c - 1);
			cout << res << "\n";
			res = 0;
		}
	}
	return 0;
}