수열과 쿼리 21 [BOJ 16975]

레이지 세그를 구현할 수 있는지를 테스트하는 문제이다.

 

문제 요약

길이 10만 이하의 수열 A에서 다음 2가지 쿼리를 처리한다.

1 i j k : 수열 A의 구간 [i, j]에 k를 더한다.

2 x : A[x]를 출력한다.

 

사용 알고리즘

Lazy propagation을 이용한 세그먼트 트리

 

풀이

Lazy propagation을 이용한 세그먼트 트리를 그대로 구현해주면 된다. 나도 처음에는 구현하기 힘들었는데, 구현이 힘들어도 너무 좌절하지 말고 계속 구현해보는 것을 추천한다. 초반에는 다른 사람의 코드를 거의 따라치다시피 해도 좋으니 정상적으로 작동하는 코드를 완성하는 것에 의의를 두는 것이 좋다. 한 번 이해하기 시작하면 활용도가 매우 높아 PS할 때 많은 도움이 된다.

 

코드

#include<bits/stdc++.h>
#define SIZE 100010
using namespace std;
typedef long long ll;
ll N,M,A[SIZE],seg[SIZE<<2],lz[SIZE<<2];
ll init(int l,int r,int node)
{
    if(l==r)return seg[node]=A[l];
    int mid=l+r>>1;
    return seg[node]=init(l,mid,node<<1)+init(mid+1,r,(node<<1)+1);
}
void prop(int l,int r,int node)
{
    seg[node]+=(r-l+1)*lz[node];
    if(l^r){
        lz[node<<1]+=lz[node];
        lz[(node<<1)+1]+=lz[node];
    }
    lz[node]=0;
}
void update(int l,int r,int node,int s,int e,ll w)
{
    prop(l,r,node);
    if(l>e||r<s)return;
    if(s<=l&&r<=e){
        lz[node]+=w;
        return;
    }
    int mid=l+r>>1;
    update(l,mid,node<<1,s,e,w);
    update(mid+1,r,(node<<1)+1,s,e,w);
}
ll query(int l,int r,int node,int s)
{
    prop(l,r,node);
    if(l>s||r<s)return 0;
    if(l==s&&r==s)return seg[node];
    int mid=l+r>>1;
    return query(l,mid,node<<1,s)+query(mid+1,r,(node<<1)+1,s);
}
int main()
{
    ios_base::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>N;
    for(int i=1;i<=N;i++)cin>>A[i];
    init(1,N,1);
    for(cin>>M;M--;){
        ll tc,s,e,w;cin>>tc>>s;
        if(tc-1)cout<<query(1,N,1,s)<<'\n';
        else{
            cin>>e>>w;
            update(1,N,1,s,e,w);
        }
    }
}