17353번: 하늘에서 떨어지는 1, 2, ..., R-L+1개의 별 - 세그먼트 트리(lazy)

lazy propagation으로 구간 갱신을 하더라도 각 배열마다 서로 다른 수를 갱신해야 하므로, 다른 방법이 필요하다.

별들의 수에 대한 차이 배열을 생각해 보자. 공차가 1인 등차수열을 더한다면, 차이 배열에서 구간 \([L, R]\)에 1을 더하고, \(R+1\)에 \(-(R-L+1)\)을 더하는 것으로 나타낼 수 있다.

 

따라서 차이 배열을 구현한 후, 위의 방법을 lazy propagation으로 구현하면 된다.

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef tuple<int, int, int> tiii;

const int MAXN = 100001;
int N, Q;
ll A[MAXN];
ll segTree[MAXN*4], lazy[MAXN*4];

void setLazy(int s, int e, int idx)
{
    ll val = lazy[idx];
    lazy[idx] = 0;
    
    segTree[idx] += (e-s+1)*val;
    if(s != e)
    {
        lazy[idx*2] += val;
        lazy[idx*2+1] += val;
    }
}

void update(int s, int e, int l, int r, int idx, ll val)
{
    if(lazy[idx]) setLazy(s, e, idx);
    
    if(e < l || r < s) return;
    if(l <= s && e <= r)
    {
        segTree[idx] += (e-s+1)*val;
        if(s != e)
        {
            lazy[idx*2] += val;
            lazy[idx*2+1] += val;
        }
        return;
    }
    
    update(s, (s+e)/2, l, r, idx*2, val);
    update((s+e)/2+1, e, l, r, idx*2+1, val);
    segTree[idx] = segTree[idx*2] + segTree[idx*2+1];
}

ll getVal(int s, int e, int l, int r, int idx)
{
    if(lazy[idx]) setLazy(s, e, idx);
    
    if(e < l || r < s) return 0;
    if(l <= s && e <= r) return segTree[idx];
    return getVal(s, (s+e)/2, l, r, idx*2) + getVal((s+e)/2+1, e, l, r, idx*2+1);
}

int main()
{
	ios::sync_with_stdio(0);
	cin.tie(0), cout.tie(0);
	
    cin >> N;
    for(int i=1; i<=N; i++) 
    {
        cin >> A[i];
        update(1, N, i, i, 1, A[i]-A[i-1]);
    }
    
    cin >> Q;
    while(Q--)
    {
        int q; cin >> q;
        if(q == 1)
        {
            int L, R; cin >> L >> R;
            update(1, N, L, R, 1, 1);
            update(1, N, R+1, R+1, 1, -(R-L+1));
        }
        else 
        {
            int X; cin >> X;
            cout << getVal(1, N, 1, X, 1) << "\n";
        }
    }
}