본문 바로가기

ps(알고리즘)

구간 합 구하기2

 

 

구간 합 문제는 세그먼트 트리로 풀 수 있었는데 이 문제는 update를 하나의 값만 하는게 아니라 구간을 정해서 구간 안의 모든 수를 갱신한다.

그래서 원래 구현했던 대로 update를 처리하면 구간의 길이만큼 K번 해야되서 시간이 안된다.

 

이 문제는 lazy propogation 알고리즘만 알면 구현으로 끝나는 간단한 문제이다.

세그트리에 lazy라는 새로운 값을 부여하는데, 구간 업데이트를 할 때 현재 탐색하는 구간이 업데이트하려는 구간에 포함된다면,

아래의 노드들을 탐색하는 것이 아니라 그냥 lazy 배열에 값을 추가해준다.

그리고 나중에 업데이트를 하거나 쿼리를 탐색할 때 lazy 배열에 값이 존재한다면, 해당 노드의 seg 값을 구간의 길이*lazy값 만큼

업데이트를 해주고, 아래 노드 2개로 가서 각각 lazy 값을 또 추가해준다.

이런 식으로 하면 특정 노드를 필요로 해서 탐색할 때만 실질적으로 값을 변경해주면 되고, 건드릴 필요 없는 노드는 하위 노드들을

전부 다 갱신할 필요 없기 때문에 시간복잡도를 줄일 수 있다.

 

long long build(int n, int s, int e)
{
    if(s>e) return 0;
    if(s==e)
    {
        return seg[n] = arr[s];
    }
    int m = (s+e)/2;
    return seg[n] = build(n*2,s,m) + build(n*2+1,m+1,e);
}

build는 일반적인 세그트리와 동일하다.

void lazy_spread(int n, int s, int e)
{
    seg[n] += (e-s+1)*lazy[n];
    if(s!=e)
    {
        lazy[n*2] += lazy[n];
        lazy[n*2+1] += lazy[n];
    }
    lazy[n]=0;
    return;
}

lazy_spread가 특정 노드를 탐색하는데 lazy 값이 있을 때 하위 노드 2개에 lazy 값을 추가해주는 과정이다.

s==e 이면, 하위 노드가 없기 때문에 추가하지 않고, lazy 값을 seg[n]에 반영해줬기 때문에 lazy[n]은 다시 0으로 바꿔준다.

void update(int n, int s, int e, int fs, int fe, long long val)
{
    if(lazy[n]!=0)
    {
        lazy_spread(n,s,e);
    }
    if(fs<=s && e<=fe)
    {
        lazy[n] += val;
        lazy_spread(n,s,e);
        return;
    }
    if(fs>e || fe<s)
    {
        return;
    }
    int m = (s+e)/2;
    update(n*2,s,m,fs,fe,val);
    update(n*2+1,m+1,e,fs,fe,val);
    seg[n] = seg[n*2] + seg[n*2+1];
}

update 함수도 거의 동일한데, lazy 값이 있을 때 lazy_spread로 lazy 값을 반영해주고, 구간이 포함될 때는 하위 노드로 

탐색을 이어가는 것이 아니라 그냥 lazy값만 추가해준다는 점이 다르다.

 

long long query(int n, int s, int e, int fs, int fe)
{
    //printf("%d %d %d %d %d\n",n,s,e,fs,fe);
    if(lazy[n]!=0)
    {
        lazy_spread(n,s,e);
    }
    if(fs<=s && e<=fe)
    {
        return seg[n];
    }
    if(fs>e || fe<s)
    {
        return 0;
    }
    int m = (s+e)/2;
    return query(n*2,s,m,fs,fe) + query(n*2+1,m+1,e,fs,fe);
}

마지막으로 쿼리는 lazy 값이 있을 때 한번 반영을 해주고 계산을 시작한다는 것 말고는 일반적인 세그트리와 동일하다.

 

#include<stdio.h>
int N,M,K;
int cnt;
long long arr[1000004];
long long seg[3000004];
long long lazy[3000004];
long long ans[1000004];
long long build(int n, int s, int e)
{
    if(s>e) return 0;
    if(s==e)
    {
        return seg[n] = arr[s];
    }
    int m = (s+e)/2;
    return seg[n] = build(n*2,s,m) + build(n*2+1,m+1,e);
}
void lazy_spread(int n, int s, int e)
{
    seg[n] += (e-s+1)*lazy[n];
    if(s!=e)
    {
        lazy[n*2] += lazy[n];
        lazy[n*2+1] += lazy[n];
    }
    lazy[n]=0;
    return;
}
void update(int n, int s, int e, int fs, int fe, long long val)
{
    if(lazy[n]!=0)
    {
        lazy_spread(n,s,e);
    }
    if(fs<=s && e<=fe)
    {
        lazy[n] += val;
        lazy_spread(n,s,e);
        return;
    }
    if(fs>e || fe<s)
    {
        return;
    }
    int m = (s+e)/2;
    update(n*2,s,m,fs,fe,val);
    update(n*2+1,m+1,e,fs,fe,val);
    seg[n] = seg[n*2] + seg[n*2+1];
}
long long query(int n, int s, int e, int fs, int fe)
{
    //printf("%d %d %d %d %d\n",n,s,e,fs,fe);
    if(lazy[n]!=0)
    {
        lazy_spread(n,s,e);
    }
    if(fs<=s && e<=fe)
    {
        return seg[n];
    }
    if(fs>e || fe<s)
    {
        return 0;
    }
    int m = (s+e)/2;
    return query(n*2,s,m,fs,fe) + query(n*2+1,m+1,e,fs,fe);
}
int main()
{
    scanf("%d %d %d",&N,&M,&K);
    for(int i=1;i<=N;i++)
    {
        scanf("%lld",&arr[i]);
    }
    build(1,1,N);
    for(int i=1;i<=M+K;i++)
    {
        int a;
        scanf("%d",&a);
        if(a==1)
        {
            long long b,c,d;
            scanf("%lld %lld %lld",&b,&c,&d);
            update(1,1,N,b,c,d);
        }
        if(a==2)
        {
            long long b,c;
            scanf("%lld %lld",&b,&c);
            cnt++;
            ans[cnt] = query(1,1,N,b,c);
        }
    }
    for(int i=1;i<=K;i++)
    {
        printf("%lld\n",ans[i]);
    }
}

전체 코드이다.

'ps(알고리즘)' 카테고리의 다른 글

책 구매하기2 - 최대유량(에드몬드-카프 알고리즘)  (1) 2025.03.10
K번째 최단경로 찾기  (0) 2025.03.02
경찰차  (0) 2025.02.28
박성원  (0) 2025.02.26
발전소  (0) 2025.02.26