[웰노운] Lazy Segment Tree

서론

Lazy Segment Tree에 대해 알아보자.

일반적인 Segment Tree는 update가 특정 점에 대해서만 가능했다.

하지만 전체적인 구간에 대하여 update를 한다면 어떨까?

당연히 특정 구간의 모든 점에 대하여 update를 해줄 수 있겠지만.

구간이 크다면, 즉 구간의 모든 점에 대해 업데이트를 해준다면 문제가 생기기에, 이를 해결해야한다.

본론

문제점

우선 이전의 점 업데이트가 왜 $O(log n)$인지 생각을 해봐야 할 것 같다.

직관적인 이유는, 이진 트리 구조이므로 어떤 노드에서는 자식 중 하나를 고를 것이며, 이진 트리의 높이는 $\log n$이기 때문이다.

근데 우리가 구간 업데이트를 구현하려고 하면, 다음과 같은 문제점이 존재한다.

우리가 트리의 모든 노드에 대해 업데이트를 진행하려고 하면, 시간 복잡도가 $O(\log n)$ 초과로 폭발할 것이다.

어떻게 해야 시간 복잡도를 $O(\log n)$으로 만들 수 있을까?

아이디어

어떠한 구간에 대해서, 노드를 3가지 경우로 분리할 수 있다고 생각하자.

1. 구간이 하나도 겹치지 않는 경우
2. 구간이 완전히 겹치는 경우
3. 구간이 걸치는 경우

가장 문제가 되는 부분은 완전히 겹치는 경우이다. 이 경우를 해결할 수 있을까?

조금 생각을 해보면, 완전히 겹치는 경우라면 그 아래 노드에 가해지는 연산도 같으니.

그 구간의 자식 노드들이 방문할 때에 추가적인 연산을 해주는 lazy[i]라는 추가적인 공간을 만들어 주면 된다.

업데이트가 log n인 간단한 증명

핵심은 각 트리 레벨(깊이)마다 쿼리 구간과 경계에서 만나는 노드는 많아야 2개라는 것이다.

앞서 말했듯이 구간이 완전히 겹치는 경우를 해결했기 때문에, 우리가 탐색하는 노드는 구간이 걸치는 경우에 한해서이다.

이는 상수 계수이므로 시간복잡도는 $O(\log n)$이다.

lazy[i]의 정의

lazy[i]를 정의하는 방법은 2개가 존재한다.

lazy[i]를 노드가 방문했을 때 처리해야 할 연산으로 정의할 수도 있고.

lazy[i]를 i 노드가 구간에 걸쳐서, 자식 노드로 전파해야 할 때 수행해야 할 연산으로 정의할 수도 있다.

전자가 맨 처음 세그트리를 이해하는 데는 쉬운 구현이나, 후자는 CP, PS에서 널리 쓰이는 방식이므로 후자를 기준으로 작성한다.

전파

이러한 전파를 다음과 같은 시점에서 수행하면 된다.

어떠한 노드가 자식 노드로 전파해야 할 때, lazy[i]를 알맞게 전파한다.

따라서 update 시점과 query 시점에서 업데이트 push()를 진행하면 된다.

의사 코드

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

class lazySegTree {
   public:
    int n;
    vector<ll> a, seg, lazy;

    lazySegTree(vector<ll> v) : n(v.size()), a(v), seg(4 * n), lazy(4 * n) { init(1, 1, n); }

    ll 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 push(int node, int l, int r) {
        if (lazy[node] == 0 || l == r) return; // 리프 노드면, return.
        int mid = (l + r) / 2;
        seg[2 * node] += (mid - l + 1) * lazy[node];
        seg[2 * node + 1] += (r - mid) * lazy[node];
        lazy[2 * node] += lazy[node];
        lazy[2 * node + 1] += lazy[node];
        lazy[node] = 0;
    } // 핵심 전파함수

    void update(int node, int l, int r, int ql, int qr, ll k) {
        if (r < ql || qr < l) return; // 1. 구간이 겹치지 않음.

        if (ql <= l && r <= qr) {
            lazy[node] += k;
            seg[node] += (r - l + 1) * k;
            return;
        } // 2. 구간이 완전히 겹침.

        push(node, l, r); // 3. 구간이 걸쳐서, 전파.
        int mid = (l + r) / 2;
        update(node * 2, l, mid, ql, qr, k);
        update(node * 2 + 1, mid + 1, r, ql, qr, k);
        seg[node] = seg[node * 2] + seg[node * 2 + 1];
    }

    ll query(int node, int l, int r, int k) {
        if (l == r) return seg[node];
        push(node, l, r);
        int mid = (l + r) / 2;
        if (k <= mid)
            return query(node * 2, l, mid, k);
        else
            return query(node * 2 + 1, mid + 1, r, k);
    }
};