백준 2042 구간 합 구하기 c++ (세그먼트 트리)

문제 출처 : https://www.acmicpc.net/problem/2042

2042번: 구간 합 구하기

문제

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

입력

첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

출력

첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

알고리즘 분류

• 자료 구조
• 세그먼트 트리

풀이

세그먼트 트리 기본 문제이다.

오늘 세그먼트 트리를 처음으로 접하였는데, 처음엔 dp로 문제를 해결하려다가 dp로는 풀 수 없을 것 같아서 찾아보니 세그먼트 트리라는 자료구조를 이용해야 했다.

세그먼트 트리란 배열의 부분 합을 구할 때 사용하는 개념으로,

일반 배열에서 어떠한 값을 바꾸는 연산은 O(1), a부터 b까지의 합을 구하는 연산은 O(N)이다.

위의 두 연산이 m번 주어질 때, 두 연산을 진행하는데 걸리는 시간 복잡도는 O(NM)이다.

이러한 경우 세그먼트 트리는 O(logN)만큼의 시간 복잡도로 굉장히 효율적으로 해결할 수 있다.

 

주어진 입력에서 배열의 값은 범위가 따로 나와있지 않지만, 정답은 모두 long long 타입을 초과하지 않는다.

따라서 세그먼트 트리와 입력 값을 저장할 배열은 모두 long long 타입으로 선언해야 한다.

 

이 문제는 세그먼트 트리 기본 문제이므로, 세그먼트 트리에 대한 공부를 하면 바로 풀 수 있기에 풀이는 생략한다.

대신, 세그먼트 트리 자료구조에 대한 포스팅을 따로 할 예정인데, 이번 하반기 공채를 받아치느라 알고리즘, 자료구조에 대한 포스팅은 조금 시간이 걸릴 것 같다.

 

코드

#include <iostream>
#include <vector>
#include <math.h>

#define MAX 1000000

using namespace std;
int n, m, k;
//1<=n<=1000000
//1<=m,k<=10000
long long arr[MAX];
long long makeSegmentTree(vector<long long> &segmentTree, int node, int start, int end) {
	if (start == end) {
		return segmentTree[node] = arr[start];
	}
	int mid = (start + end) / 2;

	return segmentTree[node] = makeSegmentTree(segmentTree, node * 2, start, mid) + makeSegmentTree(segmentTree, node * 2 + 1, mid+1, end);
	
}

void updateSegmentTree(vector<long long> &segmentTree, int node, int start, int end,int idx, long long diff) {
	if (idx<start || idx > end) return;
	segmentTree[node] += diff;
	if (start != end) {
		int mid = (start + end) / 2;
		updateSegmentTree(segmentTree, node * 2, start, mid, idx, diff);
		updateSegmentTree(segmentTree, node * 2 + 1, mid + 1, end, idx, diff);
	}
}
long long sumSegmentTree(vector<long long> &segmentTree, int node, int left, int right, int start, int end) {
	if (left>end || right<start) return 0;
	if (left <= start && right >= end) return segmentTree[node];
	int mid = (start + end) / 2;
	return sumSegmentTree(segmentTree, node * 2, left, right, start, mid) + sumSegmentTree(segmentTree, node * 2 + 1, left, right, mid+1, end);
}
int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	cin >> n >> m >> k;
	for (int i = 0; i < n; i++) {
		cin >> arr[i];
	}
	int treeDepth = ceil(log2(n));
	int treeSize = 1 << (treeDepth + 1);
	vector<long long> segmentTree(treeSize);
	
	makeSegmentTree(segmentTree,1,0,n-1);
	for (int i = 0; i < m + k; i++) {
		int order, left;
		long long right;
		cin >> order>> left >> right;
		if (order == 1) {
			//change
			updateSegmentTree(segmentTree, 1, 0, n - 1, left - 1, right-arr[left - 1]);
			arr[left - 1] = right;
		}
		else {
			//sum
			cout << sumSegmentTree(segmentTree, 1, left-1, right-1, 0, n - 1)<<'\n';
		}
	}
	return 0;
}