백준 12837: 가계부 (Hard)

문제 요약

작성될 가계부 프로그램은 두 가지 동작을 처리해야 한다. 첫 번째는 월곡이의 생후 p일에 수입/지출 내용을 추가하는 것이다. 수입은 양수, 지출은 음수의 형태로 입력이 들어온다. 두 번째는 월곡이의 생후 p일부터 q일까지 잔고가 변화한 값을 구하고 출력하는 것이다. 월곡이가 빚을 지고 있을 수도 있기에 어떤 i에 대해서 생후 i일의 잔고는 음수일 수 있다.

문제 분류

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

문제 풀이

일정 구간의 합을 빠르게 구하는 세그먼트 트리를 이용하였다. 문제에 따라서 처음 상태는 0이므로 따로 초기화할 필요가 없다. 대신, update하는 연산만 O(logn)으로 맞춰주면 되는데, update가 필요한 부분은 갱신해주고, 그렇지 않은 부분은 갱신을 스킵한다.

현재 방문중인 노드 번호와 갱신할 인덱스 번호가 존재할 때, 현재 방문중인 노드의 구간이 해당 인덱스 번호를 포함한다면, 해당 노드 역시 갱신이 필요하므로 자식 노드에 대한 합으로 갱신시키며 해당 값을 리턴한다. 그렇지 않고 변경할 인덱스가 현재 구간을 벗어난 경우, 자기 자신을 리턴하여 자신에 대한 갱신을 무시한다. 이렇게 해야만 갱신 시 시간초과가 뜨지 않는다. 마지막으로 리프노드에 도착하면, 해당 노드가 인덱스일 경우 갱신해준다. 단, 기존의 값을 무시하는 것이 아니라 누적임에 유의한다.

풀이 코드

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

using namespace std;

long long n, seg[4000000];

long long update(int l, int r, int loc, int idx, int val)
{
	if (idx < l || idx > r) return seg[loc];
	if (l == r) {
		if(l == idx) seg[loc] += val;
		return seg[loc];
	}
	int mid = (l + r) / 2;

	seg[loc] = update(l, mid, loc * 2 + 1, idx, val) + update(mid + 1, r, loc * 2 + 2, idx, val);
	return seg[loc];
}

long long sum(int L, int R, int l, int r, int idx)
{
	if (l > R || r < L) return 0;
	if (L <= l && R >= r) return seg[idx];
	int mid = (l + r) / 2;
	return sum(L, R, l, mid, idx * 2 + 1) + sum(L, R, mid + 1, r, idx * 2 + 2);
}


int main()
{
	int q, in1, in2, in3;
	scanf("%lld%d", &n, &q);

	for (int i = 0; i < q; i++) {
		scanf("%d%d%d", &in1, &in2, &in3);
		in2--;
		if (in1 == 1)
			update(0, n - 1, 0, in2, in3);
		else
			printf("%lld\n", sum(in2, in3 - 1, 0, n - 1, 0));
	}

	return 0;
}