[알고리즘] 구간합을 위한 세그먼트 트리와 펜윅 트리의 구현

 구간합을 반복적으로 빠르게 계산하기 위해서 세그먼트 트리와 펜윅 트리를 이용한다.

세그먼트 트리와 펜윅 트리는 원래의 데이터 값을 저장하는 방법이 아니라, 이름에서 알 수 있듯 데이터를 구간별로 저장하는 형식을 가진다.  간단하게 말하여 부분합을 위하여 i번째까지의 데이터를 모두 합한 값을 유지하고 있다고 보면 된다.

 세그먼트 트리는 리프노드에는 원래의 데이터를 그리고 그 부모 노드는 아래의 리프 노드들의 합을 저장하고 있다고 보면된다. 이러한 특징을 바탕으로 특정 구간의 값을 찾는 행위는 O(logN) 그리고 값을 변경 하는 행위 또한 O(lgN)만에 수행할 수 있다.

 세그먼트 트리의 경우 필요 메모리량이 복잡하게 계산된다.

Full binary tree의 경우 N이 2의 제곱꼴이면 2*N-1개의 노드가 필요하며, 만약 제곱꼴이 아니면 높이 ceil(logN)에 대하여 2^(H+1)-1개의 노드가 필요하게 된다. 따라서 여유롭게 계산해서 4N개의 노드가 필요하게 된다.

자세한 설명은 백준 세그먼트 트리를 참고.

 

 펜윅 트리의 경우 원래의 데이터를 유지하는 경우 중복 저장되는 문제를 해결하기 위하여 이미 정해진 구간마다 합을 업데이트 한다. 펜윅 트리는 바이너리 인덱스 트리라고도 한다.

 자세한 설명은 백준 펜윅 트리를 참고.

 

세그먼트 트리 

lld d[1000005];
lld T[4000005];
int n,m,k;
lld build(int l, int r, int id){
  if(l == r) return T[id] = d[l];
  int mid = l+r; mid/=2;
  return T[id] = build(l,mid,id*2) + build(mid+1,r,id*2+1);
}
lld query(int x, int y, int l, int r, int id){
  if(x > r || y < l) return 0;
  if(x <= l && y >= r) return T[id];
  int mid = l + r;  mid/=2;
  return query(x,y,l,mid,id*2)+query(x,y,mid+1,r,id*2+1);
}
void update(int x, lld v, int l, int r, int id){
  if(l==r){
    T[id]=v;
    return;
  }
  int mid = (l+r)/2;
  if(x <= mid)
    update(x, v, l, mid, 2*id);
  else
    update(x, v, mid+1, r, 2*id+1);
  T[id] = T[id*2]+T[id*2+1];
}

build(1,n,1);

 

 

펜윅 트리 

int h[1000005];
lld T[1000005];
int n, m, k;
lld rsum(lld id){
    lld ret =0;
    while(id){
        ret += T[id];
        // remove lsb
        id &= (id-1);
    }
    return ret;
}
void update(lld id, lld v){
    while(id <= n){
        T[id]+=v;
        // id&-id find lsb
        id+=(id&-id);
    }
}

lld query(lld x, lld y){
    return rsum(y)- rsum(x-1);
}