Segment tree with Lazy propagation

1. Segment tree

Segment tree를 설명하기 가장 좋은 문제가 구간의 최솟값 찾기(RMQ)라고 생각한다.

(세그먼트 트리가 아닌 다른 방식으로 RMQ를 구하는 더 좋은 풀이(복잡도는 비슷하나 더 깔끔하고 상수배 빠른)를 작성한 적이 있다. 그 글에서 사용했던 parse table과 Segment tree가 개념적으로 엄청나게 비슷하다.)

하지만 RMQ가 아닌 Segment tree(또는 BIT)로만 풀어야 하는 경우가 존재한다. 바로 배열이 변경되는 경우이다. RMQ의 경우에는 배열이 변경되지 않는다는 가정 아래에서 전처리(O(nlog(n)))를 해주었는데 Segment tree의 경우 하나의 원소 변경(O(log(n)))이 가능하다.(물론 탐색에 쓰이는 복잡도는 O(log(n)))으로 같다.

배열이 변경되는 경우 구간의 최솟값을 구하는 세그먼트 트리를 기준으로 알고리즘과 코드를 설명하겠다.

segment tree는 자식 노드가 2개 또는 0개있는 트리 모양을 갖는다. 1번 노드는 루트 노드로써 arr[1]~arr[n]의 최솟값을 저장한다.

이제 1번 노드는 2번, 3번 노드로 나누어진다. (모든 세그 트리에서 x번 노드는 2x,2x+!번 노드로 나누어진다.) 2번 노드는 arr[1]-arr[n/2]의 최솟값, 3번 노드는 arr[n/2]-arr[n]의 최솟값 이런 식으로 구간을 반 씩 나눠 갖는다.

따라서 우리가 segment tree에서 값을 변경할 때는 x번째 원소를 포함하는 모든 구간들을 바꾸어 주어야 한다. 이는 1번 노드부터 트리 탐색을 해나가면서 해당 노드의 range가 x를 포함하면 값을 변경해주고, 아니라면 탐색을 종료하면 된다.

segment tree의 최대 높이가 log(n)이므로 하나의 원소를 변경하는데 필요한 시간복잡도는 O(log(n))이다.

구간 [a,b]에서의 최솟값을 찾는 쿼리에서는 RMQ에서 했던 방식을 생각해보면 쉽다. 루트에서 트리를 탐색하면서 어떤 노드의 range가 [a,b]에 속한다면 탐색을 더 할 필요 없이 해당 노드의 값이 최솟값의 후보가 된다. 노드의 range가 [a,b]와 겹친다면 탐색을 계속 하고, 겹치지 않으면 더 탐색할 필요가 없다.

아래는 https://www.acmicpc.net/problem/10868 (최솟값)의 정답 코드이다.

#include<bits/stdc++.h>
using namespace std;
const int SZ=100005;
const int INF=INT_MAX;
int arr[SZ];
int seg[4*SZ];  //segment tree의 크기는 4*data size 정도면 충분하다.
void seg_update(int idx, int l, int r, int g, int x){   
  if(g<l || r<g) return;
  seg[idx]=min(seg[idx],x);
  if(l==r) return;
  seg_update(2*idx, l, (l+r)/2, g, x);
  seg_update(2*idx+1, (l+r)/2+1, r, g, x);
}
int seg_find(int idx, int l, int r, int fl, int fr){
  if(r<fl || fr<l) return INF;
  if(fl<=l && r<=fr) return seg[idx];
  return min(seg_find(2*idx,l,(l+r)/2,fl,fr),seg_find(2*idx+1,(l+r)/2+1,r,fl,fr));
}
int main(void){
	ios::sync_with_stdio(false); cin.tie(NULL);
	int n,m; cin>>n>>m;
	for(int i=1;i<4*SZ;i++) seg[i]=INF;
	for(int i=1;i<=n;i++){
		int x; cin>>x;
		seg_update(1,1,n,i,x);
	}
	for(int i=0;i<m;i++){
		int x,y; cin>>x>>y;
		cout<<seg_find(1,1,n,x,y)<<'\n';
	}
	return 0;
}

2. Segment tree ver.2(range update, point find)

2020-08-12 업데이트

1에서 소개한 일반적인 세그먼트 트리는 (포인트 업데이트, 구간 검색)을 하는 자료구조이지만, 그 반대(구간 업데이트, 포인트 검색)도 특정 조건에서는(최댓값, 최솟값) 비슷한 원리로 구현이 가능하다.

업데이트 함수는 1의 find에서 했던 것처럼 해당 구간에 속하는 모든 segment들을 갱신한다.

검색 함수에서는, 어떤 포인트 x에 대하여 x가 포함될 수 있는 구간을 모두 확인하는 방법으로 최댓값(또는 최솟값)을 찾을 수 있다.

구간 [a,b]에 최댓값을 갱신하고 seg[x]값을 찾는 코드는 다음과 같다.

class Segment_Tree{
private:
    vi seg;
    int n;
public:
    Segment_Tree(int n){
        this->n=n;
        seg.resize(4*n,0);
    }
    void update(int idx, int l, int r, int fl, int fr, int g){
        if(l>fr || r<fl) return;
        if(fl<=l && r<=fr){
            seg[idx]=max(seg[idx],g);
            return;
        }
        int m=(l+r)/2;
        update(2*idx,l,m,fl,fr,g);
        update(2*idx+1,m+1,r,fl,fr,g);
    }
    int find(int idx, int l, int r, int f){
        if(f<l || f>r) return 0;
        if(l==r) return seg[idx];
        int m=(l+r)/2;
        return max(seg[idx],max(find(2*idx,l,m,f),find(2*idx+1,m+1,r,f)));
    }
};

3. Segment tree with lazy propagation

range update+ range find

기존 세그먼트 트리는 update과정에서 하나의 값만 변경되었다면 lazy propagation은 update과정에서 하나 이상의 구간에 대하여 값을 변경하여도 update와 find의 복잡도가 쿼리당 log(n)이 나오게 된다.

lazy propagation에 대한 기본적인 아이디어는 다음과 같다. update를 lazy하게(게으르게) 하는 것이다. 조금 더 자세히 설명하자면

즉, 우리는 3개의 함수로 Segment Tree with Lazy propagation을 보기 좋게 구현할 수 있다. 배열에서 구간을 update하거나 구간의 합을 구하는 백준: 구간 합 구하기2 문제를 기반으로 lazy propagation을 설명하겠다.

첫번째 함수는 lazy를 update하는 함수이다.

void update_lazy(int idx, int l, int r){
  seg[idx]+=(r-l+1)*lazy[idx];  // 구간의 합을 구하는 문제이므로 구간의 크기만큼
  if(l!=r){ // lazy tree의 자손이 있는 경우
    lazy[2*idx]+=lazy[idx];
    lazy[2*idx+1]+=lazy[idx];
  }
  lazy[idx]=0;  //다 썼으니 초기화
}

위에서부터 쌓여있는 lazy들을 탐색 과정에서 자식들에게 물려주고, 자기 자신은 0로 초기화하는 간단한 함수이다.

두번째 함수는 segment tree를 update하는 함수이다.

void seg_update(int idx, int l, int r, int fl, int fr, int dif){
  update_lazy(idx, l, r);
  if(r<fl || fr<l) return;
  if(fl<=l && r<=fr){
    seg[idx]+=(r-l+1)*dif;
    if(l!=r){ // 현재 노드는 해결했고 자식들 lazy는 줘야지.
      lazy[2*idx]+=dif;
      lazy[2*idx+1]+=dif;
    }
    return;
  }
  seg_update(2*idx,l,(l+r)/2,fl,fr,dif);
  seg_update(2*idx+1,(l+r)/2+1,r,fl,fr,dif);
  seg[idx]=seg[2*idx]+seg[2*idx+1];
}

탐색을 하고 segment tree를 update해준다. 탐색 종료 직전 노드부터 update해주면서 update한 값을 segment tree에 저장한다.

세번째 함수는 구간의 합을 구하는 find 함수이다.

int seg_find(int idx, int l, int r, int fl, int fr){
  update_lazy(idx, l, r);
  if(r<fl || fr<l) return 0;
  if(fl<=l && r<=fr)  return seg[idx];
  return seg_find(2*idx, l, (l+r)/2, fl, fr)+seg_find(2*idx+1,(l+r)/2+1,r,fl,fr);
}

백준 10999번 구간 합 구하기 2https://www.acmicpc.net/problem/10999 정답 코드이다.

위 함수들을 그대로 사용하였지만 자료형의 차이가 있어 정답 코드를 다시 한 번 올린다.

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int SZ=1000005;  
ll seg[4*SZ];  
ll lazy[4*SZ];
void upd_lazy(int idx, int l, int r){
  seg[idx]+=(ll)(r-l+1)*lazy[idx];
  if(l!=r){
    lazy[2*idx]+=lazy[idx];
    lazy[2*idx+1]+=lazy[idx];
  }
  lazy[idx]=0LL;
}
void seg_update(int idx, int l, int r, int fl, int fr, ll plus){
  upd_lazy(idx,l,r);
  if(r<fl || fr<l) return;
  if(fl<=l && r<=fr){
    seg[idx]+=(ll)(r-l+1)*plus;
    if(l!=r){
      lazy[2*idx]+=plus;
      lazy[2*idx+1]+=plus;
    }
    return;
  }
  seg_update(2*idx,l,(l+r)/2,fl,fr,plus);
  seg_update(2*idx+1,(l+r)/2+1,r,fl,fr,plus);
  seg[idx]=seg[2*idx]+seg[2*idx+1];
}
ll seg_find(int idx, int l, int r, int fl, int fr){
  upd_lazy(idx,l,r);
  if(r<fl || fr<l) return 0LL;
  if(fl<=l && r<=fr) return seg[idx];
  return seg_find(2*idx,l,(l+r)/2,fl,fr)+seg_find(2*idx+1,(l+r)/2+1,r,fl,fr);
}
int main(void){
   ios::sync_with_stdio(false); cin.tie(NULL);
   int n,m,k; cin>>n>>m>>k;
   for(int i=1;i<=n;i++){
      int x; cin>>x;
      seg_update(1,1,n,i,i,x);
   }
   for(int i=0;i<m+k;i++){
      int q; cin>>q;
      if(q==1){
         int x,y,z; cin>>x>>y>>z;
         if(x>y) swap(x,y);
         seg_update(1,1,n,x,y,z);
      }
      else{
         int x,y; cin>>x>>y;
         if(x>y) swap(x,y);
         cout<<seg_find(1,1,n,x,y)<<'\n';
      }
   }
   return 0;
}