PS/Template

SegmentTree (Lazy propagation)

surkim 2024. 10. 10. 16:21
  1. 세그먼트에서 구간을 빠르게 변경하고 싶을 때
class SegmentTree {
public:
    SegmentTree(int size): n(size) {
        tree.assign(4 * n, 0);
        lazy.assign(4 * n, 0);
    }

    SegmentTree(const vector<long long>& arr): n(arr.size()) {
        tree.assign(4 * n, 0);
        lazy.assign(4 * n, 0);
        build(arr, 1, 0, n - 1);
    }

    void range_update(int left, int right, long long value) {
        update_util(1, 0, n - 1, left, right, value);
    }

    long long query(int left, int right) {
        return query_util(1, 0, n - 1, left, right);
    }

private:
    int n;
    vector<long long> tree;
    vector<long long> lazy;

    void build(const vector<long long>& arr, int node, int start, int end) {
        if (start == end) {
            tree[node] = arr[start];
        } else {
            int mid = (start + end) / 2;
            build(arr, node * 2, start, mid);
            build(arr, node * 2 + 1, mid + 1, end);
            tree[node] = tree[node * 2] + tree[node * 2 + 1];
        }
    }

    void propagate(int node, int start, int end) {
        if (lazy[node] != 0) {
            tree[node] += (end - start + 1) * lazy[node]; 
            if (start != end) {
                lazy[node * 2] += lazy[node]; 
                lazy[node * 2 + 1] += lazy[node];
            }
            lazy[node] = 0;
        }
    }

    void update_util(int node, int start, int end, int left, int right, long long value) {
        propagate(node, start, end);

        if (right < start || end < left) {
            return;
        }

        if (left <= start && end <= right) {
            tree[node] += (end - start + 1) * value;
            if (start != end) {
                lazy[node * 2] += value;
                lazy[node * 2 + 1] += value;
            }
            return;
        }

        int mid = (start + end) / 2;
        update_util(node * 2, start, mid, left, right, value);
        update_util(node * 2 + 1, mid + 1, end, left, right, value);
        tree[node] = tree[node * 2] + tree[node * 2 + 1];
    }

    long long query_util(int node, int start, int end, int left, int right) {
        propagate(node, start, end);

        if (right < start || end < left) {
            return 0;
        }
        if (left <= start && end <= right) {
            return tree[node];
        }
        int mid = (start + end) / 2;
        long long left_result = query_util(node * 2, start, mid, left, right);
        long long right_result = query_util(node * 2 + 1, mid + 1, end, left, right);

        return left_result + right_result;
    }
};

'PS > Template' 카테고리의 다른 글

MergeSortTree  (0) 2024.10.15
SegmentTree (구간 합)  (0) 2024.10.10
Template index  (0) 2024.10.10