Segment Tree

  Segment Tree는 어떤 배열에서 구간에 관한 정보를 저장해 놓는 자료구조이다. 누적합 배열을 사용할 때에는 $O(N)$으로 비교적 빠르긴 하지만 배열의 원소를 교체할 때에도 $O(N)$이므로 원소를 교체하는 쿼리가 많이 주어지는 경우에는 시간초과가 나기 쉽다.

 

 반면에 Segment Tree는 이진트리의 형태의 자료구조이므로, 원소 교체나 구간합을 구할 때의 시간복잡도가 $O(logN)$으로 매우 빠르다.

 

1. Segment Tree의 구조

 

 위 그림은 {1,9,3,8,4,5,5,9,10,3,4,5}라는 배열에 대한 구간합을 저장한 'Segment Tree'를 나타낸 것이다. 이진트리 형식으로 구성되어있으며 어떤 노드에 적혀 있는 숫자는 자식 노드들의 적혀 있는 숫자들의 합이라는 것을 알 수 있다. 또한 리프노드들에 적힌 숫자는 원래 배열의 숫자이다.

 

왼쪽에 있는 큰 서브트리

 

 자세히 보면 트리에서 왼쪽에 있는 큰 서브트리는 배열의 0~5까지의 인덱스에 대한 Segment Tree임을 알 수 있다. 이런식으로 Segment Tree는 재귀적으로 정의되어있다.

 

 

2. Segment Tree의 구현

(1) Init

  Segment Tree를 초기화해주는 함수이다. 구현은 재귀적으로 하며, Segment Tree에서 노드들의 값들은 Tree라는 배열에 저장해준다.

int arr[NMAX],tree[4*NMAX];

int init(int s, int e, int node){
    if(s==e) return tree[node] = arr[s];
    int mid = (s+e)/2;
    return tree[node] = init(s,mid,node*2) + init(mid + 1,e,node*2 + 1);
}

 

 $arr[]$는 원래 배열을 의미하고 $tree[]$는 Segment Tree를 의미한다. $s$와 $e$는 범위에 관한 정보를 담는 인덱스들이다. $s$는 시작하는 인덱스, $e$는 끝나는 인덱스이다. $s==e$인 경우는 범위가 1, 즉 리프노드를 의미하므로 해당 노드 번호에 $arr$의 값을 저장해준다. 참고로 $tree$의 메모리가 $arr$의 메모리의 4배인 이유는 $N=2^{k}$일 때 필요한 $tree$의 노드수가 $2^{k+1}-1$로 거의 2배이고, 2의 거듭제곱이 아닐 때에는 2배보다 크며 4배보다는 작기 때문에 충분하게 4배로 한 것이다.

 

 

(2) Sum

 Segment Tree를 이용해서 구간합을 구해주는 함수이다. 이 또한 재귀적으로 구현해준다.

int sum(int s, int e, int node, int l, int r){
    if(l>e || r<s) return 0; //범위 밖인 경우
    if(l<=s&&e<=r) return tree[node]; //범위와 완전히 일치하는 경우
    int mid = (s+e)/2;
    return sum(s,mid,node*2,l,r) + sum(mid+1,e,node*2+1,l,r);
}

 여기서 $l$과 $r$은 구간합을 구하고자하는 구간의 시작 인덱스와 끝 인덱스를 의마한다. 

 

 

(3) Update

 원소를 교체할 때 Segment Tree를 업데이트 해주는 함수이다. 리프노드를 교체한 뒤, 그 리프노드에 영향을 받는 모든 노드들에 값을 더해주거나 빼서 업데이트 해주는 방식이다.

arr[idx] = arr[idx] + val;

void update(int s,int e,int node,int idx,int val){
    if(idx<s||idx>e) return;
    tree[node] += val;
    if(s==e) return;
    int mid = (s+e)/2;
    update(s,mid,node*2,idx,val);
    update(mid+1,e,node*2+1,idx,val);
}

 여기서 $idx$는 교체하고자 하는 인덱스를 의미하고, $val$은 원래 값과 교체하려는 값의 차이이다. 주의해야할 점은 $arr[idx]$의 값도 같이 업데이트를 해주어야한다는 점이다.