[ BOJ ] 17353 : 하늘에서 떨어지는 1, 2, ..., R-L+1개의 별 ( PLATINUM 2 ) / C

문제

욱제의 은밀한 취미 중 하나는 매일 밤하늘을 감상하는 것이다. 😓 욱제는 하늘의 별들이 다음과 같은 규칙들을 따르며 떨어지는 걸 관찰했다.

1. 별이 떨어지는 위치는 N개의 점이다. 점은 순서대로 1, 2, ..., N의 번호를 갖는다.
2. 매일 밤 별들은 1, 2, ..., N의 연속한 부분 구간 [L, R]에 떨어진다.
3. [L, R]에 별이 떨어지면, 각 점에는 순서대로 1, 2, ..., R-L+1개의 별이 떨어진다. 다시 말해, L에는 1개, L+1에는 2개, ..., R에는 R-L+1개의 별이 떨어진다.

욱제는 하늘에서 떨어지는 별들을 기록하다가 잠이 들어버렸다!! 혹시나 했지만 역시나, 여러분은 욱제를 대신해 아래의 쿼리를 수행해야 한다. (ㅎㅎ;; ㅈㅅ.. ㅋㅋ!!)

• 1 L R: [L, R]에 별이 떨어진다. (1 ≤ L ≤ R ≤ N)
• 2 X: 점 X에 떨어진 별의 개수의 합을 출력한다. (1 ≤ X ≤ N)

입력

첫째 줄에 별이 떨어지는 점의 개수 N이 주어진다. (1 ≤ N ≤ 105)

둘째 줄에 욱제가 잠들기 전까지 세어 놓은, 이미 떨어진 별들의 개수 A1, ..., AN이 공백을 사이에 두고 주어진다. (0 ≤ A1, ..., AN ≤ 106)

셋째 줄에는 쿼리의 개수 Q가 주어진다. (1 ≤ Q ≤ 105)

넷째 줄부터 Q개의 줄에는 쿼리가 한 줄에 하나씩 주어진다.

출력

2번 쿼리에 대한 답을 한 줄에 하나씩 출력한다.

풀이 과정

정말 풀이가 며칠 동안 생각이 안나다가 이곳저곳 도움받고 하다보니 지하철에서 맞은 문제..

처음으로 세그먼트 트리에 계차수열을 도입했던 문제였다.

기존에는 원본 값을 트리에 저장하고 합 공식을 이용해 증가시키려고 했지만..

lazy에 값을 어떻게 분배할지가 문제여서 막막했었다.

이때 tree에 원본 값이 아닌 앞 원소와의 차를 저장했다.

트리에 원본 값을 저장하게 되면 그때그때 쿼리를 수행하는 것이 아닌 값만 먼저 많이 증가시켜 놓을 때 lazy 값이 이상해지는 문제를 가지고 있었다. 한 요소에 +3을 저장했다가 +9를 저장했다가 +1 을 저장해버리면 lazy 값을 child 노드로 분배할 때 어떻게 해야할지 명확한 기준을 세울 수 없었다. ( 노드마다 큐를 만들까 했지만 무모한 발상.. )

그런데, 원본 값이 아닌 앞 원소와의 차를 tree에 저장하게 된다면, 실제 증가되는 값은 +1, +2, +3 ... 이라서 그런지 tree에 1씩만 증가시켜 줘도 성공적으로 값을 바꿀 수 있었다. +1, +2, +3 ... 을 증가시키는 과정에서, tree에 앞 원소와의 차가 저장되어 있기에 앞 원소와의 차이를 1씩만 증가시켜도 되는 것이다. ( 3 1 -1 수열에 1 2 3을 증가시키면 4 3 2 (4, -1, -1), 5 5 5 (5, 0, 0), 6 7 8 ( 6, 1, 1 ).. 이런 식으로 앞 원소와의 차가 동일한 간격(1)으로 양수 방향으로 커짐. ) 

변화하는 값 다음의 마지막 원소만 Update 해주면 그대로 앞 원소와의 차이를 잘 보여주는 계차수열이기에, 이 원리를 그대로 이용해서 코드에 적용시켰다.

#include <stdio.h>

int arr[100001];
long long int tree[400004];
long long int lazy[400004];

long long int init(int start, int end, int idx) {
	lazy[idx] = 0;	    
	if (start == end) {
        if (start == 0) tree[idx] = arr[start];
        else tree[idx] = arr[start] - arr[start-1];
		return tree[idx];
	}
	int mid = (start + end) / 2;
	tree[idx] = init(start, mid, idx * 2) + init(mid + 1, end, idx * 2 + 1);
	return tree[idx];
}

void update_lazy(int start, int end, int idx) {
	if (lazy[idx] != 0) {
		tree[idx] += (end-start+1)*lazy[idx];
#ifdef DEBUG
		printf("-- [%d - %d] value %lld increase (lazy)\n", start, end, lazy[idx]);
#endif
		if (start != end) {
			lazy[idx * 2] += lazy[idx];
			lazy[idx * 2 + 1] += lazy[idx];
		}
		lazy[idx] = 0;
	}
}

void update_range(int start, int end, int idx, int left, int right) {
	int mid = (start + end) / 2;
	update_lazy(start, end, idx);
	if (end < left || right < start) return;
	if (left <= start && end <= right) {
		tree[idx] += (end-start)+1;
#ifdef DEBUG
		printf("-- [%d - %d] value %lld increase\n", start, end, end-start+1);
#endif
		if (start != end) {
			lazy[idx * 2] += 1;
			lazy[idx * 2 + 1] += 1;
		}
		return;
	}
	update_range(start, mid, idx * 2, left, right);
	update_range(mid + 1, end, idx * 2 + 1, left, right);
	tree[idx] = tree[idx * 2] + tree[idx * 2 + 1];
}

long long int sums(int start, int end, int idx, int left, int right) {
	update_lazy(start, end, idx);
	if (end < left || right < start) return 0;
	if (left <= start && end <= right) return tree[idx];

	int mid = (start + end) / 2;
	return sums(start, mid, idx * 2, left, right) + sums(mid + 1, end, idx * 2 + 1, left, right);
}

#ifdef DEBUG
void tree_debug(int start, int end, int idx) {
	if (start == end) {
		printf("[%d]\t\t %lld ( lazy : %lld )\n", start, tree[idx], lazy[idx]);
		return;
	}
	else {
		printf("[%d-%d]\t %lld ( lazy : %lld )\n", start, end, tree[idx], lazy[idx]);
		int mid = (start + end) / 2;
		tree_debug(start, mid, idx * 2);
		tree_debug(mid + 1, end, idx * 2 + 1);
	}
}
#endif

void update_final(int start, int end, int idx, int what, int val) {
    update_lazy(start, end, idx);
    if(what < start || end < what) return;
    if(start == end) {
        tree[idx] -= val;
        return;
    }

    int mid = (start + end) / 2;
    update_final(start, mid, idx * 2, what, val);
    update_final(mid+1, end, idx * 2 + 1, what, val);
    tree[idx] = tree[idx*2] + tree[idx*2+1];
}

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

	init(0, n - 1, 1);

	int q;
	scanf("%d", &q);
	
	for (int i = 0; i < q; i++) {
#ifdef DEBUG
		printf("\n");
		tree_debug(0, n - 1, 1);
		printf("\n");
#endif
		int mode;
		scanf("%d", &mode);

		if (mode == 1) {
			int n1, n2;
			scanf("%d %d", &n1, &n2);
			update_range(0, n - 1, 1, n1 - 1, n2 - 1);
            if (n2 != n) update_final(0, n-1, 1, n2, n2-n1+1);
		}
		else {
			int n1;
			scanf("%d", &n1);
			printf("%lld\n", sums(0, n - 1, 1, 0, n1 - 1));
		}
	}

#ifdef DEBUG
	printf("\n");
	tree_debug(0, n - 1, 1);
	printf("\n");
#endif
	return 0;
}

정말 어려웠지만 풀고나니 재밌었던 세그먼트 트리 문제였다.