Algoteka
Log in to submit your own samples!

Lazy Propagation Segment Tree

Problem by oml1111
# Tech tags Title Creator Created date
1 0
2022-08-12 03:28
View all samples for this language (1 verified and 0 unverified)

Readable, templated implementation | C++ |

By oml1111 |
0 likes

A templated implementation of the range sum Segment Tree with Lazy Propagation. Provides \(O(\log N)\) range assignment, range addition and range query operations.

Code

#include<vector>

//T needs functions add, multiply, combine and operator()
template<class T>
struct LazySegmentTree {
    struct Node {  // Structure for the nodes of the Lazy Segment Tree
        Node *lch, *rch;
        T value;
        T *lazy_add, *lazy_set;  
        int size;
        
        void set_mark(T val) {
            if(lazy_add != NULL) delete lazy_add;
            lazy_add = NULL;
            if(lazy_set != NULL) delete lazy_set;
            lazy_set = new T(val);
        }
        void add_mark(T val) {
            if(lazy_add != NULL) *lazy_add = lazy_add->add(val);
            else lazy_add = new T(val);
        }
        
        void propagate() {  // Apply the marked operations and propagate them to children
            if(lazy_set) {
                T to_set = *lazy_set;
                value = to_set.multiply(size);
                if(lch != NULL && rch != NULL) {
                    lch->set_mark(to_set);
                    rch->set_mark(to_set);
                }
                delete lazy_set;
                lazy_set = NULL;
            }
            if(lazy_add) {
                T toAdd = *lazy_add;
                value = value.add(toAdd.multiply(size));
                if(lch != NULL && rch != NULL) {
                    lch->add_mark(toAdd);
                    rch->add_mark(toAdd);
                }
                delete lazy_add;
                lazy_add = NULL;
            }
        }
        
        T get_value() {  // Get the value of the node in a way that ensures all operations have been applied
            propagate();
            return value;
        }
        
        void update() {  // Update the value of the node based on the value of children
            if(lch != NULL) {
                value = lch->get_value().combine(rch->get_value());
            }
        }
    };
    std::vector<Node> nodes;
    int n;
    
    LazySegmentTree(int nsize, std::vector<T> init = std::vector<T>()) { 
        // Initialize the lazy segment tree as a perfect binary tree large enough to contain nsize elements
        
        n = 1;
        while(n < nsize) {
            n *= 2;
        }
        nodes.resize(2*n-1, {NULL, NULL, T(), NULL, NULL, 1});
        
        for(int i=0;i<n;i++) {
            if(i < (int)init.size()) {
                nodes[n-1+i].value = init[i];  // Assign the leaf values
            }
        }
        for(int i=n-2;i>=0;i--) {
            nodes[i].value = nodes[2*i+1].value.combine(nodes[2*i+2].value);
            nodes[i].lch = &nodes[2*i+1];  // Set the left child of the noode
            nodes[i].rch = &nodes[2*i+2];  // Set the right child of the node
            nodes[i].size = nodes[i].lch->size + nodes[i].rch->size;
        }
    }
    
    void add(int l, int r, T to_add, Node* cur = NULL, int lb = 0, int rb = -1) { 
        // Adds to_add to all elements in range [l, r)
        // [lb, rb) represents the range of the current node
        
        if(r <= l) return;
        if(cur == NULL) {
            cur = &nodes[0];
            rb = n;
        }
        cur->propagate();
        
        if(l <= lb && rb <= r) {  // Node range is fully within the operated segment
            cur->add_mark(to_add);
            return;
        }
        int mid = (lb + rb) / 2;
        if(l < mid)
            add(l, r, to_add, cur->lch, lb, mid);  // Recursive call on left subtree
        if(mid < r)
            add(l, r, to_add, cur->rch, mid, rb);  // Recursive call on right subtree
        cur->update();
    }
    
    void set(int l, int r, T to_set, Node* cur = NULL, int lb = 0, int rb = -1) {
        // Sets all elements in range [l, r) to have value to_set
        // [lb, rb) represents the range of the current node
        
        if(r <= l) return;
        if(cur == NULL) {
            cur = &nodes[0];
            rb = n;
        }
        cur->propagate();
        
        if(l <= lb && rb <= r) {  // Node range is fully within the operated segment
            cur->set_mark(to_set);
            return;
        }
        int mid = (lb + rb) / 2;
        if(l < mid)
            set(l, r, to_set, cur->lch, lb, mid);  // Recursive call on left subtree
        if(mid < r)
            set(l, r, to_set, cur->rch, mid, rb);  // Recursive call on right subtree
        cur->update();
    }
    
    T get(int l, int r, Node* cur = NULL, int lb = 0, int rb = -1) {
        // Returns the result of applying `combine` on all elements in range [l, r)
        // [lb, rb) represents the range of the current node
        
        if(r <= l) return T();
        if(cur == NULL) {
            cur = &nodes[0];
            rb = n;
        }
        cur->propagate();
        
        if(l <= lb && rb <= r)  // Node range is fully within the queried segment
            return cur->get_value();
        
        std::vector<T> ret;
        int mid = (lb + rb) / 2;
        if(l < mid)
            ret.push_back(get(l, r, cur->lch, lb, mid) );  // Recursive call on left subtree
        if(mid < r)
            ret.push_back(get(l, r, cur->rch, mid, rb) );  // Recursive call on right subtree
        
        if(ret.size() == 2) return ret[0].combine(ret[1]);
        return ret[0];
    }
};

struct StreeSumInt {  // Integer container allowing for interval sum queries on the Lazy Segment Tree
    int val;
    StreeSumInt(int nval = 0) {val = nval;}
    StreeSumInt add(StreeSumInt r) {
        return val + r.val;
    }
    StreeSumInt multiply(int size) {
        return val * size;
    }
    StreeSumInt combine(StreeSumInt r) {
        return val + r.val;
    }
};

typedef LazySegmentTree<StreeSumInt> LazySegmentTreeIntSum;

int main() {
    LazySegmentTreeIntSum stree(10);  // Initializes Lazy Segment Tree covering 10 elements
    stree.set(3, 7, 5);  // Set elements in positions 3 ... 6 to have value 5
    stree.add(5, 8, 3);  // Add 3 to elements in positions 5 ... 7
    int range_sum = stree.get(4, 9).val;  // Get the sum of elements in positions 4 ... 8
}

Further reading

Data Structures, Advanced Concepts: Lazy Propagation - courses.cs.ut.ee
Segment Tree: Range updates (Lazy Propagation) - cp-algorithms.com

References

classes
std::vector en.cppreference.com cplusplus.com
functions
std::vector::resize en.cppreference.com cplusplus.com
std::vector::size en.cppreference.com cplusplus.com
std::vector::vector en.cppreference.com cplusplus.com

Problem Description

Implement a Segment Tree data structure with Lazy Propagation for tracking range sums. Set all elements in some range to some value, then increment all values in some range, then output the sum of all elements in some range.

View sample discussion (0 comments)
View problem discussion (0 comments)