[백준 10999] 구간 합 구하기 2 (C++)

10999번: 구간 합 구하기 2

세그먼트 트리, Lazy propagation

---

처음에 그냥 세그먼트 트리를 사용했다.

구간을 update할 때마다 리프노드까지 들어가서 update하고, 되돌아오면서 update값을 더해서 tree에 적용해줬다.

 

제출하면서도 매우 찝찝했다. 값을 수정하는 구간이 1~N까지라면 값을 update하는데 N의 시간이 걸리게 될 것이다..

구간합을 찾는 쿼리는 log(N)이라서 둘을 더하면 O(NM + Klog(N))이 되어 NM때문에 시간초과가 날거같았고 당연히 TLEㅠ

 

알고리즘 분류를 먼저 보니까 느리게 갱신되는 세그먼트 트리 라고 하는 분류가 있었다.

Lazy segment tree 를 어렴풋이 들은적이있어서 구글에서 배우기 위해 검색했다.

 

블로그 3개쯤 보면서 대충 감을 잡아가는 중에 아래 글을 우연히 봤다!

세그먼트 트리 나중에 업데이트 해야지!

 

lazy segment tree에 대해서 완벽히 설명된 글인듯 하다. 바로 이해할 수 있었고 코드까지 나와있어서 아주 편하게 이해했다.

어서 응용문제를 풀고싶다!

 

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
#include <cstring>
#define fasti ios_base::sync_with_stdio(false); cin.tie(0);
#define fastio ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#define INF 1e9+7
#define pii pair<int, int>
#define MAX 1000001
typedef long long ll;
// typedef pair<int, int> pii;
 
using namespace std;
 
int N, M, K;
ll arr[MAX];
ll tree[MAX * 4];
ll lazy[MAX * 4];
 
void input(){
    cin >> N >> M >> K;
    for(int i = 1; i <= N; i++){
        cin >> arr[i];
    }
}
 
ll make_tree(int startidx, int endidx, int node){
    if(startidx == endidx) return tree[node] = arr[startidx];
    int mid = (startidx + endidx) / 2;
    return tree[node] = make_tree(startidx, mid, node*2) +\
                        make_tree(mid+1, endidx, node*2+1);
}
 
void update_lazy(int startidx, int endidx, int node){
    if(lazy[node]){
        tree[node] += (lazy[node] * (endidx - startidx +1));
        if(startidx != endidx){
            lazy[node*2] += lazy[node];
            lazy[node*2+1] += lazy[node];
        }
        lazy[node] = 0;
    }
}
 
ll getquery(int startidx, int endidx, int lidx, int ridx, int node){
    update_lazy(startidx, endidx, node);
    if(ridx < startidx || lidx > endidx) return 0;
    if(lidx <= startidx && endidx <= ridx) return tree[node];
    int mid = (startidx + endidx) / 2;
    return getquery(startidx, mid, lidx, ridx, node*2) + \
           getquery(mid+1, endidx, lidx, ridx, node*2+1);
}
 
void update(int startidx, int endidx, int lidx, int ridx, ll up, int node){
    // 미루다가 필요할 때 update한다
    update_lazy(startidx, endidx, node);
    if(ridx < startidx || lidx > endidx) return;
    if(lidx <= startidx && endidx <= ridx){
        tree[node] += (up * (endidx - startidx + 1));
        if(startidx != endidx){
            lazy[node*2] += up;
            lazy[node*2+1] += up;
        }
        return;
    }
    
    int mid = (startidx + endidx) / 2;
    update(startidx, mid, lidx, ridx, up, node*2);
    update(mid+1, endidx, lidx, ridx, up, node*2+1);
    tree[node] = tree[node*2] + tree[node*2+1];
}
 
void solve(){
    make_tree(1, N, 1);
    int a, b, c;
    ll d;
    for(int i = 0; i < K+M; i++){
        cin >> a;
        if(a & 1){
            cin >> b >> c >> d;
            update(1, N, b, c, d, 1);
        }
        else{
            cin >> b >> c;
            cout << getquery(1, N, b, c, 1) << "\n";
        }
    }
}
 
int main(){
    fastio
    input();
    solve();
    
    return 0;
}