Segment Tree
(content/data_structures/segment_tree.h)

• View this file on GitHub
• Last update: 2024-02-01 17:09:56-05:00

Depends on

• My starter code (content/utils/template.h)

Verified with

• tests/segment_tree.test.cpp

Code

#pragma once
#include "../utils/template.h"

/**
 * @brief Segment Tree
 */

template<class T, class L> struct segment_tree{
	int SZ;
	vector<T> val;
	function<T(T, T)> merge;
	function<T(T, L)> apply;
	
	void init(int n, T def, function<T(T, T)> f, function<T(T, L)> g){
		SZ = n;
		while(n&(n-1)) n++; // next largest power of 2
		val.resize(2*n, def);
		merge = f;
		apply = g;
	}

	#define lc (rt<<1)
	#define rc (rt<<1|1)
	#define nm ((nl+nr)/2)
	
	T _query(int l, int r, int nl, int nr, int rt){
		if(r < nl or nr < l) return val[0];
		if(l <= nl and nr <= r) return val[rt];
		return merge(_query(l, r, nl, nm, lc), _query(l, r, nm+1, nr, rc));
	}
	T query(int l, int r){ return _query(l, r, 0, SZ-1, 1); }
	
	void _update(int i, L x, int nl, int nr, int rt){
		if(nl == nr){
			val[rt] = apply(val[rt], x);
			return;
		}
		i <= nm ? _update(i, x, nl, nm, lc) : _update(i, x, nm+1, nr, rc);
		val[rt] = merge(val[lc], val[rc]);
	}
	void update(int i, L x){ _update(i, x, 0, SZ-1, 1); };

	#undef lc
	#undef rc
	#undef nm
};

#line 1 "content/utils/template.h"
/**
 * @brief My starter code
 */

#include <bits/stdc++.h>
#define all(x) (x).begin(), (x).end()
#define makeunique(x) (x).erase(unique((x).begin(), (x).end()), (x).end());

#ifdef LOCAL
template<typename T> void pr(T a){std::cerr<<a<<std::endl;}
template<typename T, typename... Args> void pr(T a, Args... args){std::cerr<<a<<' ',pr(args...);}
#else
template<typename... Args> void pr(Args... args){}
#endif

using namespace std;
using ll = long long;
using pii = pair<int, int>;
#line 3 "content/data_structures/segment_tree.h"

/**
 * @brief Segment Tree
 */

template<class T, class L> struct segment_tree{
	int SZ;
	vector<T> val;
	function<T(T, T)> merge;
	function<T(T, L)> apply;
	
	void init(int n, T def, function<T(T, T)> f, function<T(T, L)> g){
		SZ = n;
		while(n&(n-1)) n++; // next largest power of 2
		val.resize(2*n, def);
		merge = f;
		apply = g;
	}

	#define lc (rt<<1)
	#define rc (rt<<1|1)
	#define nm ((nl+nr)/2)
	
	T _query(int l, int r, int nl, int nr, int rt){
		if(r < nl or nr < l) return val[0];
		if(l <= nl and nr <= r) return val[rt];
		return merge(_query(l, r, nl, nm, lc), _query(l, r, nm+1, nr, rc));
	}
	T query(int l, int r){ return _query(l, r, 0, SZ-1, 1); }
	
	void _update(int i, L x, int nl, int nr, int rt){
		if(nl == nr){
			val[rt] = apply(val[rt], x);
			return;
		}
		i <= nm ? _update(i, x, nl, nm, lc) : _update(i, x, nm+1, nr, rc);
		val[rt] = merge(val[lc], val[rc]);
	}
	void update(int i, L x){ _update(i, x, 0, SZ-1, 1); };

	#undef lc
	#undef rc
	#undef nm
};

Back to top page