본문 바로가기
Data science/알고리즘 학습

백준 2042 : 구간 합 구하기 (C++) [Segment tree]

by ggoboogi_house 2022. 4. 1.
반응형

  기본적인 segment tree를 연습해 볼 수 있는 문제이다.

 

**문제풀이 방법

- Segment tree 생성, 새 값이 들어왔을 때 update 함수, 합을 구하는 함수 작성

***코드

#include <stdio.h>
#include <math.h>
#include <vector>
typedef long long int ll;

using namespace std;


// start, end : node가 담당하고 있는 구간
// left, right : 합을 구하려는 구간

// Segment tree 초기화
ll init(vector<ll>& arr, vector<ll>& tree, int node, int start, int end) {
	if (start == end) {
		return tree[node] = arr[start];
	}
	else {
		int mid = (start + end) / 2;
		int left_child = node * 2;			// node의 왼쪽 자식
		int right_child = node * 2 + 1;		// node의 오른쪽 자식
		return tree[node] = init(arr, tree, left_child, start, mid) + init(arr, tree, right_child, mid + 1, end);
	}
}

// Tree update
void update(vector<ll>& tree, int node, int start, int end, int idx, ll diff) {
	if (idx < start || idx > end) return;
	tree[node] = tree[node] + diff;
	if (start != end) {  // leaf 노드가 아니면,
		int mid = (start + end) / 2;
		int left_child = node * 2;			// node의 왼쪽 자식
		int right_child = node * 2 + 1;		// node의 오른쪽 자식
		update(tree, left_child, start, mid, idx, diff);
		update(tree, right_child, mid+1, end, idx, diff);
	}
}

// 구간합 구하기
ll sum(vector<ll>& tree, int node, int start, int end, int left, int right) {
	if (left > end || right < start) {  // 합을 구하려는 구간이 node와 겹치지 않는 경우
		return 0;
	}
	if (left <= start && end <= right) { // node가 포함하는 구간이 합을 구하려는 구간에 속하는 경우 해당 node의 값 반환
		return tree[node];
	}
	int mid = (start + end) / 2;
	int left_child = node * 2;			// node의 왼쪽 자식
	int right_child = node * 2 + 1;		// node의 오른쪽 자식
	return sum(tree, left_child, start, mid, left, right) + sum(tree, right_child, mid + 1, end, left, right);
}


int main() {
	
	//freopen("input.txt", "r", stdin);

	int N, M, K;
	scanf("%d %d %d", &N, &M, &K);
	vector<ll> arr(N);
	int tree_height = (int)ceil(log2(N));
	int tree_size = (1 << (tree_height+1));
	vector<ll> tree(tree_size);
	M += K;

	for (int i = 0; i < N; i++) {
		scanf("%lld", &arr[i]);
	}
	init(arr, tree, 1, 0, N - 1);

	while (M--) {
		int a;
		scanf("%d", &a);
		if (a == 1) {	// a==1이면 b번째 수를 c로 바꾸기
			int b;
			ll c;
			scanf("%d %lld", &b, &c);
			b -= 1;					
			ll diff = c - arr[b];		
			arr[b] = c;
			update(tree, 1, 0, N - 1, b, diff);
		}
		else if (a == 2) {
			int b, c;
			scanf("%d %d", &b, &c);
			b--;    // 인덱스가 0부터 시작하므로 1 빼주기
			c--;    // 인덱스가 0부터 시작하므로 1 빼주기
			printf("%lld\n", sum(tree, 1, 0, N - 1, b, c));
		}
	}

	return 0;
}

 

****주의할 점

- 인덱스가 0부터 시작하는 것 유념
- 벡터 초기화

 

<Reference>

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

- https://yabmoons.tistory.com/432

- https://ongveloper.tistory.com/252

 

반응형