Segment Tree Impl, C++ (clang)

Segment Tree Impl

//Microsoft (R) C/C++ Optimizing Compiler Version 19.00.23506 for x86

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>

using namespace std;

void buildSTUtil(vector<int>& arr, vector<int>& st, int arrStart, int arrEnd, int stPos)
{    
    if(arrStart == arrEnd)
    {
      // we have reached the leaf.
      st[stPos] = arr[arrStart];
      return;
    }
    
    int mid = arrStart + ((arrEnd - arrStart) >> 1);
    int leftRootPos = (stPos * 2) + 1;
    int rightRootPos = (stPos * 2) + 2;
    
    // construct left and right sub trees of the segment tree.
    buildSTUtil(arr, st, mid+1, arrEnd, rightRootPos);
    buildSTUtil(arr, st, arrStart, mid, leftRootPos);
    
    // fill the current node, which is the max of left and right subtree summary nodes.
    st[stPos] = max(st[leftRootPos], st[rightRootPos]);
    
    // cout << "pos: " << stPos << " left: " << leftRootPos << " rightPos: " << rightRootPos << "  st[pos]: " << st[stPos] << endl;
}

vector<int> buildST(vector<int> arr)
{
    // calculate number of levels in final tree
    int N = arr.size();
    int height = (int)((ceil)(log2(N)));
    int stSize = 2 * pow(2, height) - 1;
    
    // allocate segment tree.
    vector<int> st(stSize, 0);
    buildSTUtil(arr, st, 0, arr.size()-1, 0);    
    return st;
}

void updateSTUtil(vector<int>& arr, vector<int>& st, int arrStart, int arrEnd, int offset, int value, int stPos)
{
    if(offset < arrStart || offset > arrEnd) return;
    
    st[stPos] = max(st[stPos], value);
    
    if(arrStart == arrEnd) return;
    
    int mid = arrStart + ((arrEnd - arrStart) >> 1);
    
    updateSTUtil(arr, st, arrStart, mid, offset, value, stPos * 2 + 1);
    updateSTUtil(arr, st, mid+1, arrEnd, offset, value, stPos * 2 + 2);
}

// update array at offset with given value.
void updateST(vector<int>& arr, vector<int>& st, int offset, int value)
{
    int oldValue = arr[offset];
    arr[offset] = value;
    updateSTUtil(arr, st, 0, arr.size()-1, offset, value, 0);
}

int getMax(vector<int>& arr, vector<int>& st, int arrStart, int arrEnd, int stPos, int i, int j)
{
    // if the range completely consumes the st node, then return node's value
    if(i <= arrStart && j >= arrEnd) return st[stPos];
    
    if(j < arrStart || i > arrEnd) return INT_MIN;
    
    int mid = arrStart + ((arrEnd - arrStart) >> 1);
    
    return max(getMax(arr, st, arrStart, mid, stPos * 2 + 1 , i , j)
               , getMax(arr, st, mid+1, arrEnd, stPos * 2 + 2, i, j));
}
           
int main()
{
    vector<int> arr = {1,2,3,4};
    
    auto st = buildST(arr);
    
    // update array at offset 1 with value 6, 
    // and also update the corresponding segment tree.
    updateST(arr, st, 1, 6);
    
    for(auto& elem : st)
    {
        cout << elem << "  " << endl;
    }
    
    // get maximum in range i, j
    int maxInRange = getMax(arr, st, 0, arr.size()-1,0, 1,3);
    
    cout << "Max: " << maxInRange << endl;
}


run \| edit \| history \| help	0