Algoteka
Log in to submit your own samples!

Segment Tree

Problem by oml1111
# Tech tags Title Creator Created date
1 0
2022-08-11 08:50
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 maximum Segment Tree data structure. Allows to set value at position in \(O(\log n)\) and to query a maximum of any range in \(O(\log n)\). Nodes are stored in an array, they are indexed starting from 1 and for each node \(i\) its child nodes are in positions \(2i\) and \(2i+1\)

Code

#include<vector>
#include<algorithm>
#include<limits>

struct SegmentTreeMax {
    typedef int T;
    static const T T_def = std::numeric_limits<T>::min();
    struct Node {
        T val;
        int l, r;
    };
    int size;
    std::vector<Node> nodes;
    
    SegmentTreeMax(int nsize = 0) {  // Create a perfect binary tree large enough to contain nsize elements
        int tpow = 1;
        while(tpow < nsize) tpow *= 2;
        size = tpow;
        
        nodes.resize(tpow*2);
        for(int i=0; i<tpow; i++) {
            nodes[tpow+i] = {T_def, i, i+1};
        }
        for(int i=tpow-1; i>0; i--) {
            nodes[i] = {T_def, nodes[i*2].l, nodes[i*2+1].r};
        }
    }
    
    void set(int pos, T new_val, int node_idx = 1) {  // Updates pos (indexing from 0)
        Node &cur = nodes[node_idx];
        if(cur.r == cur.l+1 && cur.l == pos) {
            cur.val = new_val;  // Set the value at leaf
            return;
        }
        if(pos < (cur.l+cur.r)/2) {
            set(pos, new_val, node_idx*2);  // Recursive call on left subtree
        }
        if(pos >= (cur.l+cur.r)/2) {
            set(pos, new_val, node_idx*2+1);  // Recursive call on right subtree
        }
        cur.val = std::max(nodes[node_idx*2].val, nodes[node_idx*2+1].val);  // Update the value of node based on subtrees
    }
    
    T get(int l, int r, int node_idx = 1) {  // Max on [l, r)
        Node &cur = nodes[node_idx];
        if(l <= cur.l && cur.r <= r) {
            return cur.val;  // Node range is fully within query range, so just return the value on the node
        }
        T result = T_def;
        if(l < (cur.l+cur.r)/2) {
            result = std::max(result, get(l, r, node_idx*2));  // Recursive call on left subtree
        }
        if((cur.l+cur.r)/2 < r) {
            result = std::max(result, get(l, r, node_idx*2+1));  // Recursive call on right subtree
        }
        return result;
    }
};


int main() {
    SegmentTreeMax stree(10);  // Create segment tree that can contain 10 elements
    stree.set(5, 15);  // Set the value at position 5 to be 15
    int range_max = stree.get(4, 7);  // Get the sum of elements in positions 4 ... 6
}

Further reading

Data Structures: Segment Tree - courses.cs.ut.ee
Segment Tree - wikipedia.org

References

classes
std::numeric_limits en.cppreference.com cplusplus.com
std::vector en.cppreference.com cplusplus.com
functions
std::max en.cppreference.com cplusplus.com
std::numeric_limits::min en.cppreference.com
std::vector::resize en.cppreference.com cplusplus.com

Problem Description

Implement a Segment Tree data structure for tracking range maximums. Set some value at some position and then output the maximum of some range.

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