[백준] 2042번 - 구간 합 구하기 (C언어) [Gold1]

input	output
5 2 2 1 2 3 4 5 1 3 6 2 2 5 1 5 2 2 3 5	17 12

해당 문제는 배열의 인덱스 요소에서 설정한 범위내에 구간의 합을 구하는 문제이다. 일반적인 배열의 구간 합은 O(n)의 시간 복잡도를 가지만 해당 문제에서 입력으로 주어지는 모든 수는 -2^63 보다 크거나 같고 2^63 - 1보다 작거나 같은 정수이므로 시간 제한이 2초이므로 보다 빠른 알고리즘을 사용해야 한다. 

 

필자는 세그먼트 트리를 통해 해결하였는데, 세그먼트 트리는 O(logn)의 상당히 빠른 알고리즘이다. 세그먼트 트리는 몰랐던 자료구조 이기 때문에 세그먼트 트리가 무엇인지 알아보면서 풀이를 진행하겠다. 

 

세그먼트 트리는 이진 트리 구조이다. 따라서 각 노드에 대한 정보를 알 수 있다.

1. 부모노드 인덱스 * 2 = 왼쪽 자식 인덱스
2. 부모노드 인덱스 * 2 + 1 = 오른쪽 자식 인덱스 

세그먼트 트리는 정보를 받아오는 배열로 부터 리프노드를 설정한다. 배열이 {1,2,3,4,5} 라면 세그먼트 트리 형태는 아래와 같다.

                        15

                6               9

          3         3      4       5

       1    2        

리프노드는 파란색으로 표시하였다. 해당 트리에서 알 수 있는 사실은 리프 노드 1과 2의 합은 3인데, 3은 1과 2의 부모노드이다. 이어서 3과 리프노드 3의 합은 6인데 이들의 부모노드는 6이다. 이에 대한 정보들을 담고 있는 것이 세그먼트 트리이다. 

 

세그먼트 트리를 생성하는 과정은 아래와 같다.

1. 주어진 트리를 반으로 나눔
2. 나누어진 트리에서 왼쪽 트리를 재귀 호출
3. 나누어진 트리에서 오른쪽 트리를 재귀 호출
4. 1~3 과정 반복

해당 과정에 대한 소스 코드는 아래와 같다.

ll MakeSegmentTree(int start, int end, int node) { // 트리 생성
	if (start == end) // 재귀 종료 조건
		return tree[node]=arr[start];

	int mid = (start + end) / 2; // 중간값 설정
	tree[node] =  MakeSegmentTree(start, mid, node * 2) + MakeSegmentTree(mid + 1, end, node * 2 + 1); 
	// 왼쪽 트리 반환 값 + 오른쪽 트리 반환 값 저장
	return tree[node]; // 세그먼트 트리 반환
}

트리를 생성하였으니 구간 합에 대하여 알아보자. 구간 합에서 생각해볼 수 있는 경우의 수는 아래와 같다.

1. 탐색 범위가 완전히 벗어난 경우
2. 탐색 범위가 완전히 속한 경우
3. 탐색 범위가 위 두 조건을 제외하고 일부만 겹치는 경우

해당 연산에 대한 소스 코드는 아래와 같다.

ll sum(int start, int end, int node, int left, int right) { // 구간합 함수
	if (start > right || end < left) // 탐색 범위가 완전히 벗어난 경우
		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);
}

이제 세그먼트 트리 값 변경에 대해서 알아보자. 위에서 알아본 구간 합의 조건과 유사하다.

1. 변경하고자 하는 값이 탐색 범위에 완전히 속한 경우
2. 변경하고자 하는 값이 탐색 범위에 완전히 벗어난 경우

해당 연산에 대한 소스 코드는 아래와 같다.

ll UpdateSegmentTree(int start, int end, int node, int index, ll diff) {
	if (index < start || index > end) // 범위를 벗어난 경우
		return;

	tree[node] += diff; // 트리 노드 업데이트
	
	if (start == end) // 재귀 종료 조건
		return;

	int mid = (start + end) / 2;
	UpdateSegmentTree(start, mid, node * 2, index, diff);
	UpdateSegmentTree(mid + 1, end, node * 2 + 1, index, diff);
}

 

전체 소스 코드는 아래와 같다. (주석 포함)

// 부모노드번호 x 2 = 왼쪽자식노드번호, 부모노드번호 x 2 + 1 = 오른쪽자식노드번호
// 세그먼트 트리 이용하여 문제 해결
//      10		
//   3       7
//  1 2    3  4

/*
	크기가 N인 배열이 존재할 때

	1. 트리의 높이 = ceil(log2(N))

	2. 세그먼트 트리의 크기 = (1 << (트리의 높이 + 1) )

*/

// { 현재 노드 번호 , 시작 범위 , 마지막 범위 }
// 인덱스 1부터 사용

/*
1. 주어진 범위를 반으로 나눈다.

2. 나눠진 2개의 범위에 대해서 '왼쪽범위'에 대한 재귀호출을 한다.

3. 나눠진 2개의 범위에 대해서 '오른쪽범위'에 대한 재귀호출을 한다.

4. 위의 과정을 반복한다.
*/

#include <stdio.h>
#pragma warning(disable:4996)
typedef long long int ll; // long long 타입 정의
ll arr[1000000];
ll tree[4000000];

ll MakeSegmentTree(int start, int end, int node) { // 트리 생성
	if (start == end) // 재귀 종료 조건
		return tree[node]=arr[start];

	int mid = (start + end) / 2; // 중간값 설정
	tree[node] =  MakeSegmentTree(start, mid, node * 2) + MakeSegmentTree(mid + 1, end, node * 2 + 1); 
	// 왼쪽 트리 반환 값 + 오른쪽 트리 반환 값 저장
	return tree[node]; // 세그먼트 트리 반환
}

ll sum(int start, int end, int node, int left, int right) { // 구간합 함수
	if (start > right || end < left) // 탐색 범위가 완전히 벗어난 경우
		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);
}

ll UpdateSegmentTree(int start, int end, int node, int index, ll diff) {
	if (index < start || index > end) // 범위를 벗어난 경우
		return;

	tree[node] += diff; // 트리 노드 업데이트
	
	if (start == end) // 재귀 종료 조건
		return;

	int mid = (start + end) / 2;
	UpdateSegmentTree(start, mid, node * 2, index, diff);
	UpdateSegmentTree(mid + 1, end, node * 2 + 1, index, diff);
}

int main() {
	int n, x, y;
	scanf("%d %d %d", &n,&x,&y);

	for (int i = 0; i < n; i++) {
		scanf("%lld", &arr[i]);
	}

	MakeSegmentTree(0, n - 1, 1);
	int cnt = x + y;
	for (int i = 0; i < cnt; i++) {
		ll a, b, c;
		scanf("%lld %lld %lld", &a, &b, &c);
		if (a == 1) {
			UpdateSegmentTree(0, n - 1, 1, b - 1, c - arr[b - 1]);
			arr[b - 1] = c; // 배열 요소 업데이트 
		}
		else {
			printf("%lld\n", sum(0, n - 1, 1, b - 1, c - 1)); // 구간합 출력
		}
	}

	return 0;
}

해당 문제를 통해 새로운 자료구조를 접할 수 있어서 새로웠고 생각보다 문제가 까다로웠다. 

 

[2042번]

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

2042번: 구간 합 구하기

 

[github]

https://github.com/MinWook6457

MinWook6457 - Overview

[참고 사이트]

https://yabmoons.tistory.com/431

[ 세그먼트 트리(Segment Tree) ] 개념과 구현방법 (C++)