Axolotl

BOJ 16975 수열과 쿼리 21

23 Oct 2024

Segment Tree와 Lazy Propagation을 사용해 쿼리를 처리하는 문제입니다.

#include <iostream>
#include <vector>
#include <cmath> 
using namespace std; 
int N, M; 
vector<long long> arr; 
vector<long long> lazy; 
vector<long long> Tree; 
void init(vector<long long> &a, vector<long long> &t, int index, int start, int end){
    if(start == end){
        t[index] = a[start];
    }
    else{
        int mid = (start + end) / 2; 
        init(a, t, 2 * index, start, mid);
        init(a, t, 2 * index + 1, mid + 1, end); 
        t[index] = t[2 * index] + t[2 * index + 1]; 
    }
}
void update_lazy(vector<long long> &l, vector<long long> &t, int index, int start, int end){
    if(l[index] != 0){
        t[index] += (end - start + 1) * l[index];
        if(start != end){
            l[2 * index] += l[index];
            l[2 * index + 1] += l[index];
        }
        l[index] = 0; 
    }
}
void update(vector<long long> &a, vector<long long> &l, vector<long long> &t, int index, int start, int end, int left, int right, long long diff){
    update_lazy(l, t, index, start, end);
    if(left > end || start > right){
        return; 
    }
    if(left <= start && end <= right){
        t[index] += (end - start + 1) * diff; 
        if(start != end){
            l[2 * index] += diff;
            l[2 * index + 1] += diff; 
        }
        return; 
    }
    int mid = (start + end) / 2; 
    update(a, l, t, 2 * index, start, mid, left, right, diff);
    update(a, l, t, 2 * index + 1, mid + 1, end, left, right, diff);
    t[index] = t[2 * index] + t[2 * index + 1];
}
long long query(vector<long long> &l, vector<long long> &t, int index, int start, int end, int left, int right){
    update_lazy(l, t, index, start, end); 
    if(left > end || right < start){
        return 0; 
    }
    if(left <= start && right >= end){
        return t[index];
    }
    int mid = (start + end) / 2;
    long long lsum = query(l, t, 2 * index, start, mid, left, right); 
    long long rsum = query(l, t, 2 * index + 1, mid + 1, end, left, right); 
    return lsum + rsum; 
}
int main(){
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0); 
    cin >> N; 
    int height = (int)ceil(log2(N));
    int TreeSize = (1 << (height + 1));
    arr.resize(N);
    Tree.resize(TreeSize); 
    lazy.resize(TreeSize);
    for(int i = 0; i < N; i++){
        cin >> arr[i]; 
    }
    cin >> M; 
    init(arr, Tree, 1, 0 , N - 1); 
    for(int i = 0; i < M; i++){
        int flag; cin >> flag; 
        if(flag == 1){
            int b, c; long long d; cin >> b >> c >> d; 
            update(arr, lazy, Tree, 1, 0, N - 1, b - 1, c - 1, d); 
        }
        else if(flag == 2){
            int b; cin >> b; 
            cout << query(lazy, Tree, 1, 0, N - 1, b - 1, b - 1) << "\n"; 
        }
    }
}