풀이
지금까지는 업데이트를 배열 원소에 하고 구간에 대한 연산을 구하는 문제들이었다면 이 문제는 반대로 업데이트를 구간에 하고 배열 원소를 구하는 문제다. 업데이트할 때 모든 원소를 직접 리프 노드까지 업데이트하면 이 되어 시간 초과가 발생한다. 따라서 업데이트 구간에 완전히 포함되는 노드일 경우, 그 노드에만 업데이트 값을 저장한 뒤, 나중에 실제 합을 구할 때 대상 리프 노드를 포함하는 모든 부모 노드에 누적된 업데이트를 합쳐 최종 합을 구한다.
코드
#include <iostream>
#include <algorithm>
#include <vector>
#include <cmath>
using namespace std;
using ll = long long;
int n, m, k;
vector<ll> arr, seg;
// n : tree 기준 노드 번호
// s : 배열 구간 시작 인덱스
// e : 배열 구간 종료 종료 인덱스
void init(int n, int s, int e) {
if (s == e) {
seg[n] = arr[s]; // 리프 노드만 초기화
return;
}
int mid = (s + e) / 2;
init(n * 2, s, mid);
init(n * 2 + 1, mid + 1, e);
}
// n : tree 기준 노드 번호
// t : 값을 구하는 인덱스
// s : 배열 구간 시작 인덱스
// e : 배열 구간 종료 종료 인덱스
ll getsum(int n, int t, int s, int e) {
if (t < s || t > e) return 0; // 구간이 업데이트 대상 인덱스를 포함하지 0 반환
if (s == e) { // 대상 리프 노드를 찾으면
return seg[n]; // 해당 노드값 반환
}
int mid = (s + e) / 2;
ll s1 = getsum(n * 2, t, s, mid); // 왼자식 탐색
ll s2 = getsum(n * 2 + 1, t, mid + 1, e); // 오른자식 탐색
return s1 + s2 + seg[n]; // 찾은 노드 + 그동안 쌓인 업데이트 값 반환 ( 못 찾은 쪽은 0 )
}
// n : tree 기준 노드 번호
// l : 업데이트 대상 구간 시작 인덱스
// r : 업데이트 대상 구간 종료 인덱스
// p : 업데이트 값
// s : 배열 구간 시작 인덱스
// e : 배열 구간 종료 종료 인덱스
void update(int n, int l, int r, ll p, int s, int e) {
if (l <= s && e <= r) { // 구간을 완전히 포함하는 경우
seg[n] += p; // 노드 값에 업데이트 값을 누적
return;
}
if (r < s || e < l) return; // 구간이 전혀 포함되지 않는 경우, 탐색 종료
int mid = (s + e) / 2;
update(n * 2, l, r, p, s, mid); // 왼자식 탐색
update(n * 2 + 1, l, r, p, mid + 1, e); // 오른자식 탐색
}
int main()
{
ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
cin >> n; arr = vector<ll>(n + 1); seg = vector<ll>((int)pow(2, (int)ceil(log2(n)) + 1));
for (int i = 1; i <= n; i++) {
cin >> arr[i];
}
init(1, 1, n);
cin >> m;
for (int i = 1; i <= m; i++) {
int q; cin >> q;
if (q == 1) {
int a, b, c; cin >> a >> b >> c;
update(1, a, b, c, 1, n);
}
if (q == 2) {
int a; cin >> a;
cout << getsum(1, a, 1, n) << "\n";
}
}
return 0;
}