Algorithm in A..Z - Lazy Segment Tree

개념

Segment Tree

Segment Tree에서 구간의 Update를 빠르게 처리하기 위한 방법이다. Segment Tree에서 Update의 시간 복잡도는 O(logN)이기 때문에 구간의 길이가 L인 Update의 시간 복잡도는 O(LlogN)이기 때문에 쿼리의 수가 많은 경우 시간초과를 받을 수 있다.

Lazy Segment Tree는 Update를 필요한 시기에 Update하도록 미루면서 Range Update를 O(logN)으로 수행할 수 있다.

---

작동원리

구간을 Update할 경우 Segment Tree에서 참조하는 Node만 Update를 진행하고, 자식 노드의 Update는 Lazy Tree에 기록하여 Update를 미룬다.

---

시간 복잡도

init : O(NlogN)

update : O(logN)

update_range : O(logN)

query : O(logN)

---

문제

www.acmicpc.net/problem/16975

16975번: 수열과 쿼리 21

---

코드

#include <bits/stdc++.h>
using namespace std;

vector<long long> ve;
vector<long long> tree, lazy;

long long init(int index, int left, int right) {
    if (left == right) {
        return tree[index] = ve[left];
    }

    return tree[index] = init(index * 2, left, (left + right) / 2) + init(index * 2 + 1, (left + right) / 2 + 1, right);
}

void updateLazy(int index, int left, int right) {
    if (lazy[index]) {
        tree[index] += (right - left + 1) * lazy[index];

        if (left != right) {
            lazy[index * 2] += lazy[index];
            lazy[index * 2 + 1] += lazy[index];
        }

        lazy[index] = 0;
    }
}

void update(int index, int left, int right, int begin, int end, long long diff) {
    updateLazy(index, left, right);
    if (left > end || right < begin) {
        return;
    }

    if (begin <= left && right <= end) {
        lazy[index] += diff;
        updateLazy(index, left, right);
        return;
    }

    update(index * 2, left, (left + right) / 2, begin, end, diff);
    update(index * 2 + 1, (left + right) / 2 + 1, right, begin, end, diff);
    tree[index] = tree[index * 2] + tree[index * 2 + 1];
}

long long query(int index, int left, int right, int begin, int end) {
    updateLazy(index, left, right);
    if (left > end || right < begin) {
        return 0L;
    }

    if (begin <= left && right <= end) {
        return tree[index];
    }

    return query(index * 2, left, (left + right) / 2, begin, end) + query(index * 2 + 1, (left + right) / 2 + 1, right, begin, end);
}

int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);

    int n;
    cin >> n;

    ve.resize(n);
    for (auto &i : ve) {
        cin >> i;
    }

    tree.resize(1 << ((int)ceil(log2(n)) + 1));
    lazy.resize(1 << ((int)ceil(log2(n)) + 1));
    init(1, 0, n - 1);

    int k;
    cin >> k;
    while (k--) {
        int q;
        cin >> q;

        if (q == 1) {
            int begin, end, diff;
            cin >> begin >> end >> diff;

            update(1, 0, n - 1, begin - 1, end - 1, diff);
        } else if (q == 2) {
            int i;
            cin >> i;

            cout << query(1, 0, n - 1, i - 1, i - 1) << "\n";
        }
    }
}