[자료구조] 백준 16975번: 수열과 쿼리 21 (C++)

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

16975번: 수열과 쿼리 21

구간합 lazy propagation으로 구현해주었다.

>> 1번 쿼리 (구간 합 업데이트) 때문에 세그먼트 트리를 써야 했다. 그런데 구간 업데이트라서 lazy propagation이 필요했다.

lazy propagation 개념 공부할 때와 똑같은 구성이라 그런지 어렵게 느껴지진 않았었다.

#include <iostream>
#define MAX (100000 + 1)
#define ll long long
using namespace std;

ll arr[MAX];
ll sgTree[MAX * 4];
ll lazy[MAX * 4];
int n;

ll sgInit(int start, int end, int node) {
	if (start == end) {
		return sgTree[node] = arr[start];
	}
	int mid = (start + end) / 2;
	return sgTree[node] = (sgInit(start, mid, node * 2) + 
		sgInit(mid + 1, end, node * 2 + 1));
}

void lazyUpdate(int start, int end, int node) {
	if (lazy[node]) {
		sgTree[node] += (end - start + 1) * lazy[node];
		if (start != end) {
			lazy[node * 2] = lazy[node * 2] + lazy[node];
			lazy[node * 2 + 1] = lazy[node * 2 + 1] + lazy[node];
		}
		lazy[node] = 0;
	}
}

void update(int start, int end, int node, int left, int right, ll value) {
	lazyUpdate(start, end, node);
	if (left > end || right < start) return;
	if (left <= start && end <= right) {
		sgTree[node] = sgTree[node] + (end - start + 1) * value;
		if (start != end) {
			lazy[node * 2] += value;
			lazy[node * 2 + 1] += value;
		}
		return;
	}

	int mid = (start + end) / 2;
	update(start, mid, node * 2, left, right, value);
	update(mid + 1, end, node * 2 + 1, left, right, value);
	sgTree[node] = (sgTree[node * 2] + sgTree[node * 2 + 1]);
}

ll query(int start, int end, int node, int index) {
	lazyUpdate(start, end, node);
	if (index > end || index < start) return 0;
	if (start == end) {
		return sgTree[node];
	}

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

int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cin >> n;
	for (int i = 1; i <= n; i++) {
		cin >> arr[i];
	}
	sgInit(1, n, 1);
	int m;
	cin >> m;
	for (int i = 0; i < m; i++) {
		int a;
		cin >> a;
		if (a == 1) {
			int b, c;
			ll d;
			cin >> b >> c >> d;
			update(1, n, 1, b, c, d);
		}
		else {
			int b;
			cin >> b;
			cout << query(1, n, 1, b) << '\n';
		}
	}
}