# | Likes | Tech tags | Title | Creator | Created date |
---|---|---|---|---|---|
1 | 0 |
2022-08-12 03:28
|
A templated implementation of the range sum Segment Tree with Lazy Propagation. Provides \(O(\log N)\) range assignment, range addition and range query operations.
#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
}
Data Structures, Advanced Concepts: Lazy Propagation - courses.cs.ut.ee
Segment Tree: Range updates (Lazy Propagation) - cp-algorithms.com
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 |
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.