[자료구조] 세그먼트 트리

세그먼트 트리는 구간(Segment)에 대한 쿼리를 효율적으로 처리하는 트리 자료구조이다.

 

단순한 구간의 합이나, 특정 원소 업데이트를 하는 쿼리를 처리한다고 하면 누적합 배열같은 방법을 사용하는것이 편하고 쉬울 수 있다.

 

그러나 업데이트와 구간의 합을 구하는 쿼리가 몇십만개 이상이 주어진다고 하면 누적합 배열을 사용하는 경우 O(N)의 시간이 걸리기 때문에 시간이 너무 오래걸린다.

 

세그먼트트리는 이러한 상황에서 효율적으로 작동한다.

 

세그먼트 트리의 구조

세그먼트 트리는 위 모습처럼 표현되곤 한다.

상위 노드가 하위 노드의 구간을 포함하는 방식으로 구성된다.

 

세그먼트트리는

H = ⌈log₂(N)⌉ (트리의 높이) 이므로 

전체 노드 수 = 2^(H+1) - 1

 

따라서 최악의 경우 N = 2^k + 1일 때 4N에 가까운 공간복잡도가 나온다.

 

세그먼트 트리의 구현

초기화 함수

ll init(ll node, ll S, ll E) {
    if (S == E) {
        return tree[node] = initial_array[S];
    }
    ll Mid = (S + E) / 2;
    ll left_sum = init(2 * node, S, Mid);
    ll right_sum = init(2 * node + 1, Mid + 1, E);
    return tree[node] = left_sum + right_sum;
}

노드를 타고 내려가면서 재귀 호출을 하는데, 리프 노드라면 리프 노드의 값을 넣고 반환하고, 부모 노드는 양쪽 노드들의 합을 구한 후 반환한다.

 

update함수

void update(ll X, ll V, ll node, ll S, ll E) {
    if(S==E) {
        tree[node] = V;
        return;
    }
    ll Mid = (S+E)/2;
    if(X <= Mid) update(X,V,2*node,S,Mid);
    else update(X,V,2*node+1,Mid+1,E);
    tree[node] = tree[2*node] + tree[2*node+1];
}

업데이트 함수는 목적 노드에 도달할 때까지 재귀 호출을 한다.

목적 노드에 도달하면 값을 업데이트하고

다시 위로 올라가면서 변화된 값으로 부모 노드들을 바꿔준다.

 

Query함수

ll query(ll L, ll R, ll node, ll S, ll E) {
    if(R<S || E < L) {
        return 0;
    }
    if(L<=S && E<=R) {
        return tree[node];
    }
    ll Mid = (S+E)/2;
    return query(L,R,2*node,S,Mid) + query(L,R,2*node+1,Mid+1,E);
}

쿼리 함수는 L~R까지의 구간합을 구하는 함수이다.

구간을 벗어나면 0을 반환하고 구간 내에 모두 존재하면 해당 구간에 포함되는 노드의 값을 반환한다.

이런식으로 재귀 호출을 진행해서 모든 부분에 포함되는 구간의 합을 도출할 수 있다.

 

세그먼트 트리를 사용하면 업데이트, 쿼리 모두 O(logN)으로 처리할 수 있어 빠르다.