segment tree algorithm code example

Example 1: segment tree

// General Segment Tree struct
// Original source: https://codeforces.com/blog/entry/18051

template<class T> struct SegTree {

    int size;
    T defVal;
    vector<T> tree;
    T (*op)(T, T);
    
    SegTree (int size, T defVal, T (*op)(T, T))
        : size(size), defVal(defVal), tree(vector<T>(2*size)), op(op) {}
    SegTree (vector<T> v, T defVal, T (*op)(T, T))
        : SegTree(v.size()/2, defVal, op) {
        for (int i = 0; i<size; ++i) {
            tree[i+size] = v[i];
        }
        build();
    }

    void build() {

        for (int i = size-1; i>0; --i) {
            tree[i] = op(tree[i<<1], tree[i<<1|1]);
        }

    }
    T query (int l, int r) {

        l += size, r += size+1;
        T res;
        for (res = defVal; l<r; l >>= 1, r >>= 1) {
            if (l&1) res = op(res, tree[l++]);
            if (r&1) res = op(res, tree[--r]);
        }
        return res;

    }
    void update (int i, T val) {

        i += size;
        for (tree[i] = val; i>1; i >>= 1) {
            tree[i>>1] = op(tree[i], tree[i^1]);
        }

    }

};

Example 2: segmented trees

// C++ program to show segment tree operations like construction, query 
// and update 
#include <bits/stdc++.h> 
using namespace std; 

// A utility function to get the middle index from corner indexes. 
int getMid(int s, int e) { return s + (e -s)/2; } 

/* A recursive function to get the sum of values in the given range 
	of the array. The following are parameters for this function. 

	st --> Pointer to segment tree 
	si --> Index of current node in the segment tree. Initially 
			0 is passed as root is always at index 0 
	ss & se --> Starting and ending indexes of the segment represented 
				by current node, i.e., st[si] 
	qs & qe --> Starting and ending indexes of query range */
int getSumUtil(int *st, int ss, int se, int qs, int qe, int si) 
{ 
	// If segment of this node is a part of given range, then return 
	// the sum of the segment 
	if (qs <= ss && qe >= se) 
		return st[si]; 

	// If segment of this node is outside the given range 
	if (se < qs || ss > qe) 
		return 0; 

	// If a part of this segment overlaps with the given range 
	int mid = getMid(ss, se); 
	return getSumUtil(st, ss, mid, qs, qe, 2*si+1) + 
		getSumUtil(st, mid+1, se, qs, qe, 2*si+2); 
} 

/* A recursive function to update the nodes which have the given 
index in their range. The following are parameters 
	st, si, ss and se are same as getSumUtil() 
	i --> index of the element to be updated. This index is 
			in the input array. 
diff --> Value to be added to all nodes which have i in range */
void updateValueUtil(int *st, int ss, int se, int i, int diff, int si) 
{ 
	// Base Case: If the input index lies outside the range of 
	// this segment 
	if (i < ss || i > se) 
		return; 

	// If the input index is in range of this node, then update 
	// the value of the node and its children 
	st[si] = st[si] + diff; 
	if (se != ss) 
	{ 
		int mid = getMid(ss, se); 
		updateValueUtil(st, ss, mid, i, diff, 2*si + 1); 
		updateValueUtil(st, mid+1, se, i, diff, 2*si + 2); 
	} 
} 

// The function to update a value in input array and segment tree. 
// It uses updateValueUtil() to update the value in segment tree 
void updateValue(int arr[], int *st, int n, int i, int new_val) 
{ 
	// Check for erroneous input index 
	if (i < 0 || i > n-1) 
	{ 
		cout<<"Invalid Input"; 
		return; 
	} 

	// Get the difference between new value and old value 
	int diff = new_val - arr[i]; 

	// Update the value in array 
	arr[i] = new_val; 

	// Update the values of nodes in segment tree 
	updateValueUtil(st, 0, n-1, i, diff, 0); 
} 

// Return sum of elements in range from index qs (quey start) 
// to qe (query end). It mainly uses getSumUtil() 
int getSum(int *st, int n, int qs, int qe) 
{ 
	// Check for erroneous input values 
	if (qs < 0 || qe > n-1 || qs > qe) 
	{ 
		cout<<"Invalid Input"; 
		return -1; 
	} 

	return getSumUtil(st, 0, n-1, qs, qe, 0); 
} 

// A recursive function that constructs Segment Tree for array[ss..se]. 
// si is index of current node in segment tree st 
int constructSTUtil(int arr[], int ss, int se, int *st, int si) 
{ 
	// If there is one element in array, store it in current node of 
	// segment tree and return 
	if (ss == se) 
	{ 
		st[si] = arr[ss]; 
		return arr[ss]; 
	} 

	// If there are more than one elements, then recur for left and 
	// right subtrees and store the sum of values in this node 
	int mid = getMid(ss, se); 
	st[si] = constructSTUtil(arr, ss, mid, st, si*2+1) + 
			constructSTUtil(arr, mid+1, se, st, si*2+2); 
	return st[si]; 
} 

/* Function to construct segment tree from given array. This function 
allocates memory for segment tree and calls constructSTUtil() to 
fill the allocated memory */
int *constructST(int arr[], int n) 
{ 
	// Allocate memory for the segment tree 

	//Height of segment tree 
	int x = (int)(ceil(log2(n))); 

	//Maximum size of segment tree 
	int max_size = 2*(int)pow(2, x) - 1; 

	// Allocate memory 
	int *st = new int[max_size]; 

	// Fill the allocated memory st 
	constructSTUtil(arr, 0, n-1, st, 0); 

	// Return the constructed segment tree 
	return st; 
} 

// Driver program to test above functions 
int main() 
{ 
	int arr[] = {1, 3, 5, 7, 9, 11}; 
	int n = sizeof(arr)/sizeof(arr[0]); 

	// Build segment tree from given array 
	int *st = constructST(arr, n); 

	// Print sum of values in array from index 1 to 3 
	cout<<"Sum of values in given range = "<<getSum(st, n, 1, 3)<<endl; 

	// Update: set arr[1] = 10 and update corresponding 
	// segment tree nodes 
	updateValue(arr, st, n, 1, 10); 

	// Find sum after the value is updated 
	cout<<"Updated sum of values in given range = "
			<<getSum(st, n, 1, 3)<<endl; 
	return 0; 
} 
//This code is contributed by rathbhupendra

Example 3: segment tree

const int N = 1e5;  // limit for array size
int n;  // array size
int t[2 * N];

void build() {  // build the tree
  	for (int i = n - 1; i > 0; --i) t[i] = t[i<<1] + t[i<<1|1];
}

void modify(int p, int value) {  // set value at position p
  	for (t[p += n] = value; p > 1; p >>= 1) t[p>>1] = t[p] + t[p^1];
}

int query(int l, int r) {  // sum on interval [l, r)
    int res = 0;
    for (l += n, r += n; l < r; l >>= 1, r >>= 1) {
        if (l&1) res += t[l++];
        if (r&1) res += t[--r];
    }
    return res;
}

int main() {
	std::cin>>n;
    for (int i = 0; i < n; ++i) std::cin>>t[n+i];
    build();
    modify(0, 1);
    std::cout<<query(3, 11)<<'\n';
    return 0;
}

Example 4: Segment tree

void build(int node, int start, int end)
{
    if(start == end)
    {
        // Leaf node will have a single element
        tree[node] = A[start];
    }
    else
    {
        int mid = (start + end) / 2;
        // Recurse on the left child
        build(2*node, start, mid);
        // Recurse on the right child
        build(2*node+1, mid+1, end);
        // Internal node will have the sum of both of its children
        tree[node] = tree[2*node] + tree[2*node+1];
    }
}

Tags:

Misc Example