Lazy Propagation Segment Tree

wx+b 연산을 지원하는 lazy 세그 예제

 

#include <iostream>
#include <algorithm>

using namespace std;

#define TREE_WIDTH (1 << 19)

typedef long long ll;


struct linear_t {
    ll w, b;
    
    bool is_identity(void) {
        return w == 1ll && b == 0ll;
    }
};
struct node_t {
    ll sum;
    linear_t lazy;
} tree[TREE_WIDTH * 2];

linear_t compose(linear_t a, linear_t b) {
    // (x * a.w + a.b) * b.w + b.b
    return linear_t{a.w * b.w, a.b * b.w + b.b};
}


void update(int ts, int te, linear_t v, int idx = 1, int s = 1, int e = TREE_WIDTH) {

    if (e < ts || te < s)
        return;

    else if (ts <= s && e <= te) {
        
        tree[idx].sum = tree[idx].sum * v.w + v.b * (ll)(e - s + 1);
        tree[idx].lazy = compose(tree[idx].lazy, v);

        tree[idx].sum = tree[idx].sum;
        tree[idx].lazy = {tree[idx].lazy.w, tree[idx].lazy.b};

        return;
    }

    else {
        int m = (s + e) / 2;

        if (!tree[idx].lazy.is_identity()) {
            update(s, m, tree[idx].lazy, idx * 2, s, m);
            update(m + 1, e, tree[idx].lazy, idx * 2 + 1, m + 1, e);
        }

        update(ts, te, v, idx * 2, s, m);
        update(ts, te, v, idx * 2 + 1, m + 1, e);

        tree[idx].sum = tree[idx * 2].sum + tree[idx * 2 + 1].sum;
        tree[idx].lazy = {1ll, 0ll};
    }
    
    return;
}


ll query(int ts, int te, int idx = 1, int s = 1, int e = TREE_WIDTH) {

    if (e < ts || te < s)
        return 0ll;
    
    else if (ts <= s && e <= te)
        return tree[idx].sum;

    else {
        int m = (s + e) / 2;

        node_t &now = tree[idx];
        if (!now.lazy.is_identity()) {

            update(s, m, now.lazy, idx * 2, s, m);
            update(m + 1, e, now.lazy, idx * 2 + 1, m + 1, e);

            now.lazy = {1ll, 0ll};
        }

        return query(ts, te, idx * 2, s, m) + query(ts, te, idx * 2 + 1, m + 1, e);
    }
}