백준 2042: 구간 합 구하기

문제 요약

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

문제 분류

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

문제 풀이

세그먼트 트리를 이용해서 풀어야 하는 문제이다. 해당 트리를 배열로 선언하였고, 루트 노드를 0번 인덱스로 잡았다. 이에 어떠한 노드 번호를 i라고 하면, 해당 노드의 자식 노드 인덱스는 i * 2 + 1와 i * 2 + 2가 될 것이다. 또한, 해당 노드의 부모 노드 인덱스는 (i - 1) / 2가 된다.

입력받은 n에 따라 트리의 높이와 크기를 정했다. log2()와 pow()를 이용해 입력받을 트리 배열의 인덱스를 구했다. 가령 n이 5라면, log2(5)를 올림한 값 3을 구하고, 그 값을 2의 자승으로 하여 구한 값(8)에 1을 빼면 트리 리프 노드의 시작 인덱스가 된다. 조금 어렵지만, 같은 레벨의 노드는 모두 배열에서 연속되어있고, 해당 레벨의 맨 왼쪽 노드(시작 노드)의 인덱스는 해당 레벨의 노드 갯수에서 1을 뺀 값과 같다. 이를 이용한 풀이이다. ceil(log2(n))을 이용해 리프 노드의 높이를 구하고 해당 높이의 노드 개수는 당연히 2의 레벨 자승이 될 것이다. 그 값에서 1을 빼서 시작 인덱스를 구한다.

배열 크기는 넉넉하게 3백만정도 주었다. 계산상 약 2백 몇십만 나왔지만, 충분한 크기로 할당하였다.

모든 입력이 2^64 정도인 것을 확인하여 대부분의 변수를 long long으로 선언했다.

풀이 코드

#include <stdio.h>
#include <iostream>
#include <algorithm>
#include <math.h>

using namespace std;

long long seg[3000001];
long long n, m, k, idx;

void construct() {
	for (int i = idx - 1; i >= 0; i--)
		seg[i] = seg[i * 2 + 1] + seg[i * 2 + 2];
}

void update(int loc, long long val) {
	loc += idx;
	seg[loc] = val;
	for (int i = (loc - 1) / 2; i > 0; i = (i - 1) / 2)
		seg[i] = seg[i * 2 + 1] + seg[i * 2 + 2];
	seg[0] = seg[1] + seg[2];
}

long long sol(int loc, int l, int r, int nodeL, int nodeR) {
	if (r < nodeL || nodeR < l) return 0;
	if (l <= nodeL && nodeR <= r) return seg[loc];
	int mid = (nodeL + nodeR) / 2;
	return sol(loc * 2 + 1, l, r, nodeL, mid) + sol(loc * 2 + 2, l, r, mid + 1, nodeR);
}

int main()
{
	long long in1, in2, in3;
	cin >> n >> m >> k;
	idx = ceil(log2(n));
	idx = pow(2, idx) - 1;
	for (int i = idx; i < idx + n; i++)
		scanf("%lld", seg + i);
	construct();
	for (int i = 0; i < m + k; i++) {
		scanf("%lld%lld%lld", &in1, &in2, &in3);
		in2--;
		if (in1 == 1)
			update(in2, in3);
		else
			printf("%lld\n", sol(0, in2, in3 - 1, 0, idx));
	}
	return 0;
}

후일담

pow()보다는 비트 연산자 <<를 사용하여2의 제곱승을 구하는 것이 더 좋았을 것 같다.