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

김진웅·2024년 1월 17일

2042 2042번 C 백준 세그먼트 트리

baekjoon-study

목록 보기

57/59

링크
https://www.acmicpc.net/problem/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보다 작거나 같은 정수이다.

예제 입력 1

예제 출력 1

17
12

아이디어 스케치

세그먼트 트리 자료구조를 사용하는 문제이다.

문제의 예제를 바탕으로 세그먼트 트리를 그려보면

위와 같다.

이때 파란색 숫자는 배열의 인덱스이다.

위 그림에서 부모 노드와 자식노드와의 관계를 살펴보면 자식노드의 인덱스는 (부모노드의 인덱스x2) 와 (부모노드의 인덱스x2 +1)를 가진다는 것을 알 수 있다.

예제에서 1,2,3,4,5를 입력 받는데 이 값들을 리프노드에 추가해야 한다. 리프 노드에 추가하기 위해서는 리프노드의 인덱스를 구해야 하는데, 세그먼트 트리는 이진 트리로 각 depth의 시작 노드의 인덱스는 2^tepth와 같다. 즉 5개의 숫자가 주어지므로 5보다 크면서 가장 가까운 2의 지수승인 8을 리프 노드의 시작 인덱스로 가진다.

트리의 크기 = 주어진 숫자의 개수(N) $<=$ $2^k$ x 2, 즉 16을 크기로 가진다.

트리의 크기를 16으로 초기화 한 후 인덱스 8부터 8+N까지 숫자를 입력 받은 후,
tree[N] = tree[2*N] + tree[2*N+1]
위 점화식을 이용하여 부모노드의 값을 채운다.

위 과정을 통해 세그먼트 트리를 초기화한다.

다음으로는 값을 변경하는 함수를 구현해야 한다.
문제에서 주어지는 a번째 수는 입력할 때의 순서이므로 이 순서에 해당하는 인덱스를 구해야 한다.

예제에서는 첫 번째 수가 인덱스 8번에 들어갔으므로 두 번째 수는 인덱스 9번에 들어간 것을 알 수 있다. 즉 2+2^depth-1으로 나타낼 수 있다.

tree[9]에 해당하는 값을 b로 바꾸고, 각 노드의 부모노드가 부분합을 가지고 있기 때문에 수정한 노드의 부모노드를 루트노드에 도달할 때 까지 수정한다.

마지막으로 부분합을 구하는 함수를 구현해야 한다.

a번째부터 b번째까지의 부분합을 트리 자료구조를 이용하여 구하기 위해서는,

값을 변경하는 함수와 같이 a와 b에 해당하는 인덱스를 먼저 구해야 한다.

a와 b에 해당하는 인덱스를 구한 후, 각각 start_index, end_index로 정의한다.
start_index가 홀수인 경우에는 우측 자식 노드에 해당되므로 이 노드의 부모노드를 부분합에 포함하지 않고 이 노드만 부분합에 포함한다.
end_index가 짝수인 경우에는 좌측 자식 노드에 해당되므로 이 노드의 부모노드를 부분합에 포함하지 않고 이 노드만 부분합에 포함한다.

start_index의 부모노드의 인덱스를 start_index로 갱신한다. 이때 start_index/2가 아닌 (start_index+1)/2를 해야하는데, 그 이유는 우측 자식 노드에 해당되는 경우에 이 노드의 부모노드를 포함하면 안되기 때문이다.

end_index의 부모노드의 인덱스를 end_index로 갱신한다. 이때 end_index-1을 2로 나눈 값을 end_index로 설정하는데, 그 이유는 좌측 자식 노드에 해당되는 경우에는 이 노드의 부모 노드를 포함하면 안되기 때문이다.

위 그림은 위에서 설명한 로직을 나타낸 그림이다.

노란색은 부분합에 포함하는 노드로써 start_index가 9인경우 인덱스가 홀수이므로 이 노드만 부분합에 포함하고 현재 노드의 부모 노드의 다음 노드를 start_index로 설정하는 것을 볼 수 있다.
end_index 또한 인덱스가 12 ,즉 짝수이므로 이 노드만 부분합에 포함하고 현재 노드의 부모 노드의 이전 노드를 end-index로 설정하는 것을 볼 수 있다.

전체 코드

#include <stdio.h>
#include <stdlib.h>
#define ll long long

ll* tree;
int tree_size = 1;

void change(ll a, ll b);
ll partition_sum(ll start, ll end);

int main()
{
	int N, M, K;
	
	
	int command;
	ll a, b;


	scanf("%d %d %d", &N, &M, &K);
	
	// N보다 크면서 가장 가까운 2의 지수승 값을 구함
	while (tree_size < N) {
		tree_size *= 2;
	}
	

	tree = (ll*)calloc(tree_size * 2, sizeof(ll));	// 구한 2의 지수승 값에 2를 곱한 값을 크기로 하는 트리 생성

	// 입력 받는 값들을 리프노드에 저장
	for (int i = tree_size; i < tree_size + N; i++)
		scanf("%lld", &tree[i]);

	// 자식 노드의 값의 합을 부모 노드에 저장
	for (int i = tree_size - 1; i >= 1; i--)
		tree[i] = tree[2 * i] + tree[(i * 2) + 1];

	int times = M + K;	// M+K 횟수 만큼 연산
	
	while (times--) {
		scanf("%d %lld %lld",&command, &a, &b);
		if (command == 1) change(a, b);		// command가 1인 경우 값을 변경하는 함수 호출
		else if (command == 2) printf("%lld\n", partition_sum(a, b));	// command가 2인 경우 부분 합을 구하는 함수 호출
	}

	// 메모리 할당 해제
	free(tree);

	return 0;
	
}

// 값을 변경하는 함수
void change(ll a, ll b)
{
	ll target_index = a + tree_size - 1;
	ll update = b - tree[target_index];
	tree[target_index] = b;
	target_index /= 2;
	while (target_index > 0) {
		tree[target_index] += update;
		target_index /= 2;
	}
}

// 부분합을 구하는 함수
ll partition_sum(ll start, ll end) {
	ll sum = 0;
	ll start_index = start + tree_size - 1;
	ll end_index = end + tree_size - 1;

	while (start_index <= end_index) {		// start인덱스와 end인덱스가 뒤집히기 전까지 반복
		if (start_index % 2 == 1) sum += tree[start_index];		// start_index에 해당하는 노드가 우측 자식에 해당하는 경우에는 부모노드를 포함하지 않고 이 노드만 부분합에 추가
		if (end_index % 2 == 0) sum += tree[end_index];		// end_index에 해당하는 노드가 좌측 자식에 해당하는 경우에는 부모 노드를 포함하지 않고 이 노드만 부분합에 추가

		start_index = (start_index + 1) / 2;	// start_index를 부모 노드로 설정(+1을 하는 이유는 start_index가 우측 자식인 경우에는 현재 자식의 부모노드를 포함하지 않기 위함)
		end_index = (end_index - 1) / 2;		// end_index를 부모 노드로 설정(-1을 하는 이유는 end_index가 좌측 자식인 경우에는 현재 자식의 부모노드를 포함하지 않기 위함
	}

	return sum;		// 결과 반환
}