Segment Tree : 세그먼트 트리

세그먼트 트리는 특정 구간에 대해 주어지는 쿼리에 효율적으로 응답하기 위해 만들어진 자료구조이다.
(구간의 합, 특정 구간에서 가장 큰 수 / 작은 수 등)

입력 배열을 리프 노드로 하고 내부 노드는 리프 노드를 병합해서 나타낸다.
주어지는 쿼리에 대해서 병합 방식은 다를 수 있다.

세그먼트 트리를 구현하는 방법은 여러 가지가 있지만 이 글에서는 구간 합에 대해서 재귀로 구현하는 방법에 대해서 정리했다.

세그먼트 트리의 구현은 크게 3가지로 나눌 수 있다.

1. 트리 생성
2. 노드 업데이트
3. 쿼리 처리

트리 생성

세그먼트 트리의 크기(배열의 길이)는 입력 배열의 크기 n에 대해서 (int) (Math.ceil(Math.log(n) / Math.log(2))) 이다.
(아래 코드의 arr는 입력 배열, tree는 세그먼트 트리 배열을 말한다.)

// si = 세그먼트 트리 배열 인덱스 / s = 구간의 start 인덱스 / e = 구간의 end 인덱스
public int makeTree(int si, int s, int e) {
    if(s==e) return tree[si] = arr[s];
    int mid = (s+e)/2;
    return tree[si] = makeTree(si*2, s, mid) + makeTree(si*2+1, mid+1, e);
}

노드 업데이트

입력 트리에서 변경하는 요소의 인덱스를 구간의 시작과 끝 인덱스와 비교하면서 구간에 인덱스가 포함된다면 업데이트 해준다.

// si = 세그먼트 트리 배열 인덱스 / s = 구간의 start 인덱스 / e = 구간의 end 인덱스
// ai = 입력 배열에서 업데이트 하는 요소의 인덱스 / diff = 업데이트 하는 요소의 
public void update(int si, int s, int e, int ai, int diff) {
    if(ai<s || e<ai) return;
    tree[si] += diff;
    if(s==e) return;
    int mid = (s+e)/2;
    update(si*2, s, mid, ai, diff);
    update(si*2+1, mid+1, e, ai, diff);
}

쿼리 처리

쿼리로 주어지는 구간의 시작과 끝 인덱스를 현재 계산하는 구간의 시작과 끝 인덱스와 비교해준다.
1. 현재 계산하는 구간의 시작과 끝 인덱스가 쿼리로 주어지는 구간에 포함된다면 현재 탐색중인 노드에 저장된 값을 반환한다.
2. 현재 계산하는 구간의 시작과 끝 인덱스가 쿼리로 주어지는 구간에 포함되지 않는다면 0을 반환한다.
3. 위 케이스에 해당하지 않는다면 자식 노드들의 합을 리턴해준다.(재귀)

// si = 세그먼트 트리 배열 인덱스 / s = 구간의 start 인덱스 / e = 구간의 end 인덱스
// as = 쿼리 구간의 start 인덱스 / ae = 쿼리 구간의 end 인덱스
public int query(int si, int s, int e, int as, int ae) {
    if(ae<s || e<as) return 0;
    if(as<=s && e<=ae) return tree[si];
    int mid = (s+e)/2;
    return sum(si*2, s, mid, as, ae) + sum(si*2+1, mid+1, e, as, ae);
}