[웰노운] Segment Tree

서론

Segment Tree에 대해 알아보자.

본론

사실 이번 포스트는 UCPC 같은 대회에서 내 코드 템플릿을 만든 것이다.

이 구현에서 주의할 점은 다음과 같다.

1. 입력 배열은 0-based index를 따른다, 세그트리의 일반적인 구현은 1-based index이므로 주의해야 한다.
2. RSQ(구간 합 쿼리)기반으로 작성되었기 때문에 operation에 대해 수정해야한다.

의사 코드

위 내용을 구현한 c++ 코드이다. 구간 합을 업데이트하는 코드 예제이다.

class segTree {
   public:
    int n;
    vector<int> a, seg;

    segTree(int x) : n(x), a(n, 0), seg(4 * n) { init(1, 1, n); }
    segTree(vector<int> v) : n(v.size()), a(v), seg(4 * n) { init(1, 1, n); }

    int init(int node, int l, int r) {
        if (l == r) return seg[node] = a[l - 1];
        int mid = (l + r) / 2;
        return seg[node] = init(node * 2, l, mid) + init(node * 2 + 1, mid + 1, r);
    }

    void update(int node, int l, int r, int idx, int k) {
        if (idx < l || r < idx) return;
        if (l == idx && r == idx) {
            seg[node] = a[idx - 1] = k;
            return;
        }
        int mid = (l + r) / 2;
        update(node * 2, l, mid, idx, k);
        update(node * 2 + 1, mid + 1, r, idx, k);
        seg[node] = seg[node * 2] + seg[node * 2 + 1];
    }

    int query(int node, int l, int r, int ql, int qr) {
        if (qr < l || r < ql) return 0;
        if (ql <= l && r <= qr) return seg[node];
        int mid = (l + r) / 2;
        return query(node * 2, l, mid, ql, qr) + query(node * 2 + 1, mid + 1, r, ql, qr);
    }
};