세그먼트 트리

2 minute read

세그먼트 트리(Segment Tree)는 일반적으로 구간(Segment) 사이 연산의 쿼리와 수정을 빠르게 하기 위한 자료구조이다.

누적합에 대해서

세그먼트 트리는 흔히 누적합을 여러번 구할 때 많이 사용된다. 쿼리 및 수정을 항상 $O(\lg n)$만에 할 수 있기 때문이다.

그러면 세그먼트 트리의 생성과 쿼리 및 수정 방법을 알아보자.

트리의 생성은 다음과 같이 한다. (0-based index)

fn init_tree(tree: &mut Vec<i64>, arr: &Vec<i64>, i: usize, l: usize, r: usize) -> i64 {
    if l == r {
        tree[i] = arr[l];
        return arr[l];
    }
    let m = (l + r) / 2;
    tree[i] = init_tree(tree, arr, i * 2 + 1, l, m) + init_tree(tree, arr, 2 * i + 2, m + 1, r);
    return tree[i];
}

구간 쿼리는 다음과 같이 한다. ($[ql, qr)$의 원소의 누적합을 구함)

fn query(tree: &mut Vec<i64>, i: usize, l: usize, r: usize, ql: usize, qr: usize) -> i64 {
    if r < ql || l > qr {
        0
    } else if ql <= l && r <= qr{
        tree[i]
    } else {
        let m = (l + r) / 2;
        query(tree, 2 * i + 1, l, m, ql, qr) + query(tree,2 * i + 2, m + 1, r, ql, qr)
    }
}

이때 어떤 $ui$의 값을 수정하는 업데이트는 다음과 같이 한다.

fn update(tree: &mut Vec<i64>, i: usize, l: usize, r: usize, ui: usize, uv: i64) {
    if l == r {
        tree[i] = uv;
        return;
    }
    let m: usize = (l + r) / 2;
    if ui <= m {
        update(tree, 2 * i  +1, l, m, ui, uv);
    } else {
        update(tree, 2 * i + 2, m + 1, r, ui, uv);
    }
    tree[i] = tree[2 * i  +1] + tree[2 * i + 2];
}

세그먼트 트리의 속성

세그먼트 트리은 몇가지 속성을 갖는다. 일단 개수가 $n$인 배열에 대한 세그먼트 트리에 대해

공간은 약 $4n$을 차지한다.

그리고 세그먼트 트리는 무엇보다 Full Binary Tree를 가진다.

일반화

세그먼트 트리가 성립하는 연산에는 위에서 보았던 더하기 연산($+$)이 있다. 그러면 어떤 연산이 세그먼트 트리를 만들 수 있을까?

어떤 $T$에 대해서 어떤 이항연산 $\ast$에 대해서 $(T, \ast, e)$가 모노이드(Monoid)이면 연산 $\ast$은 $T$에 대한 세그먼트 트리를 만들 수 있다고 볼 수 있다. 이때 모노이드란 다음과 같은 성질을 만족하는 일종의 대수적 구조이다:

따라서 우리는 위의 구간 누적합의 경우에는 $(N_0, +, 0)$으로 위 조건을 만족하는 모노이드이므로 세그먼트 트리를 만들 수 있다고 볼 수 있었던 것이다.

그러면 Segment Tree를 어떤 모노이드라 볼 수 있는 연산($F$)에 대해서 Generic하게 만들 수도 있을 것이다.

pub trait MonoidFn<T> {
    fn call(&self, a: &T, b: &T) -> T;
    fn identity(&self) -> T;
}

struct SegTree<T: Clone, F>
where
    F: MonoidFn<T>,
{
    pub op: F,
    size: usize,
    tree: Vec<T>,
}

위와 같이 구성할 수 있다.