수열과 쿼리 39 [BOJ 19651]

수쿼 스타일의 문제를 보다보면 특정 구간에 등차수열을 더하거나, 피보나치 수열을 더하는 등의 문제를 볼 수 있다. 이 문제를 통해 등차수열에 관한 쿼리는 어떻게 처리하는지 익혀보길 바란다. 그리고 다이아라고 되어있긴 한데 왜 다이아인지 전혀 모르겠다. 플1 정도가 적당해보이고, 사람에 따라 더 내려서 생각할 수도 있을 것 같다. 푸른색 티어에 겁먹지 말고 도전해보자.

 

문제 요약

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

1 i j x y : 구간 [i, j]에 초항이 x이고 공차가 y인 등차수열 더하기

2 i j : 구간 [i, j]에 포함되는 등차수열인 구간 [l, r] 중 가장 긴 것의 길이를 출력

 

사용 알고리즘

금광 세그

 

풀이

태그에 레이지가 달려있긴 한데 사실 레이지가 필수적이지는 않다. '등차수열'의 의미를 다시 곱씹어본다면 금광 세그먼트 트리로 해결 가능하다. 등차수열이란 차가 일정한 수열이기 때문에, 인접한 수의 차를 담는 배열을 생각하자. 새로운 수열 B를 B[i]=A[i]-A[i-1]로 저장하면 아래와 같다.

그러면 같은 수가 연속되어 나오는 구간 중 가장 긴 구간을 찾는 문제로 바뀐다. 하지만 이를 구하는 것도 쉬워보이지 않는데, 여기서 조금 더 고민을 하면 된다. 만약 그 '같은 수'=x라고 알고 있었다면 금광 세그를 통해 해당 구간의 왼쪽에서부터 연속한 x의 수, 오른쪽에서부터 연속한 x의 수 등을 저장하여 해당 구간에서 최대 몇 개의 x가 연속해있는지를 관리해줄 수 있을 것이다. 그렇다면 여기서 x를 하나로 모아줄 수 없을까? 있다. 바로 A에서 B를 만들었던 것처럼, B에서 C를 똑같이 만들어주는 것이다. C[i]=B[i+1]-B[i]로 정의하자.

인접한 같은 수끼리는 빼면 0이 나오기 때문에, 결국 우리가 구하고자 하는 값은 C에서 연속한 0의 최대 개수이다. 즉, 아까 위에서 했던 금광 세그 관찰에서 x=0인 경우와 같기에 문제를 해결할 수 있다. 다르게 접근해도 같은 C 배열에 도달할 수 있다. 등차수열에서는 2*A[i]=A[i-1]+A[i+1]로 등차중항이 성립한다. 따라서 2부터 N-1까지의 i에 대해 C[i]=A[i+1]+A[i-1]-2*A[i]로 정의할 수 있는데, 계산해보면 위에서 유도한 C 배열과 같다. 이제 각 업데이트가 어떻게 되고, 금광 세그를 어떻게 처리해야하는지 생각해보자.

 

위의 A, B, C가 함께 있는 그림에서도 알 수 있듯이 구간 [s, e]가 등차수열이라면 C에서 구간 [s+1, e-1]이 0이 된다. 따라서 출력 쿼리를 처리할 때에는 C의 구간 [s+1, e-1]에서 연속한 0의 최대 개수에 2를 더하여 출력해주면 된다. 업데이트 쿼리에서 어느 지점을 업데이트 해야하는지 알기 위해 식을 써봐도 되는데, 개인적으로 그림으로 표현하는 것이 더 직관적인 것 같아서 그림으로 그려보았다. 아래 그림에서는 모든 수가 0인 수열의 일부분에 초항이 x, 공차가 y인 등차수열을 더했다.

구간 [s, e]에 더했을 경우, B에는 굉장히 많은 변화가 일어나지만 C에는 정확히 4점에만 변화가 일어난다. 이는 갱신되는 구간의 중간 부분에서의 배열 C가 변하지 않기 때문이다. 결과적으로 배열 C에서 s-1번 칸에 +x, s번 칸에 +y-x, e번 칸에 -x-(e-s+1)y, e+1번 칸에 x+(e-s)y를 더해주면 된다. 업데이트할 때에는 해당 점이 현재 0인지, 갱신 후 0인지를 바탕으로 체크해주면 일반적인 경우 쿼리 횟수를 조금 줄일 수 있는데, 필수적인 과정은 아닌 것 같다.

 

금광 세그를 관리할 때에는 다음 5가지 값을 관리해주면 된다.

int le : 왼쪽에서부터 연속한 0의 수

int ri : 오른쪽에서부터 연속한 0의 수

int mx : 연속한 0의 최대 수

int al : 구간의 길이

bool fl : 구간 내부가 모두 0이면 true, 아니면 false

(사실 fl=(mx==al)로 처리해줄 수 있어서 필수적이지는 않다.)

두 노드를 merge하는 과정은 자식 노드들의 5가지 값을 이용하면 쉽게 해결할 수 있다. 이 부분을 혼자 짤 수 있다면 이 문제에서 금광 세그가 어떻게 사용되는지 완벽히 이해한 것이다. 혼자 짜보고 아래 필자의 코드와 비교해보자.

 

이렇게 전체 시간복잡도 O(NlogN)에 문제를 해결할 수 있으며, 상수가 조금 커보이지만 N이 작고 시간제한도 2초여서 널널하게 통과된다. 필자의 코드는 272ms가 나왔다.

 

코드

#include<bits/stdc++.h>
#define SIZE 100010
using namespace std;
typedef long long ll;
struct Node
{
    ll le,ri,mx,al,fl;
};
ll N,M,A[SIZE],B[SIZE];
Node seg[SIZE<<2];
Node mer(Node x,Node y)
{
    Node ret;
    ret.fl=x.fl&y.fl;
    ret.le=x.le;
    if(x.fl)ret.le=x.al+y.le;
    ret.ri=y.ri;
    if(y.fl)ret.ri=x.ri+y.al;
    ret.al=x.al+y.al;
    ret.mx=max(max(x.mx,y.mx),x.ri+y.le);
    return ret;
}
Node init(ll l,ll r,ll idx)
{
    if(l==r){
        if(B[l]==0)return seg[idx]={1,1,1,1,1};
        return seg[idx]={0,0,0,1,0};
    }
    ll mid=l+r>>1;
    return seg[idx]=mer(init(l,mid,idx<<1),init(mid+1,r,(idx<<1)+1));
}
Node upd(ll l,ll r,ll idx,ll s,ll x)
{
    if(l>s||r<s)return seg[idx];
    if(l==s&&r==s){
        if(x==0)return seg[idx]={1,1,1,1,1};
        return seg[idx]={0,0,0,1,0};
    }
    ll mid=l+r>>1;
    return seg[idx]=mer(upd(l,mid,idx<<1,s,x),upd(mid+1,r,(idx<<1)+1,s,x));
}
Node que(ll l,ll r,ll idx,ll s,ll e)
{
    if(l>e||r<s)return {0,0,0,0,1};
    if(s<=l&&r<=e)return seg[idx];
    ll mid=l+r>>1;
    return mer(que(l,mid,idx<<1,s,e),que(mid+1,r,(idx<<1)+1,s,e));
}
void update(ll s,ll x)
{
    if(B[s]==0)upd(2,N-1,1,s,1);
    B[s]+=x;
    if(B[s]==0)upd(2,N-1,1,s,0);
}
int main()
{
    ios_base::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>N;
    for(ll i=1;i<=N;i++)cin>>A[i];
    for(ll i=2;i<N;i++)B[i]=A[i-1]+A[i+1]-2*A[i];
    init(2,N-1,1);
    for(cin>>M;M--;){
        ll tc,s,e,x,y;cin>>tc>>s>>e;
        if(tc==1){
            cin>>x>>y;
            update(s-1,x);
            update(s,-x+y);
            update(e,-x-(e-s+1)*y);
            update(e+1,x+(e-s)*y);
        }
        if(tc==2)printf("%lld\n",2+que(2,N-1,1,s+1,e-1).mx);
    }
}