Skip to content

Instantly share code, notes, and snippets.

@tanitanin
Created December 19, 2014 16:35
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save tanitanin/9e3a421486b4e168e250 to your computer and use it in GitHub Desktop.
Save tanitanin/9e3a421486b4e168e250 to your computer and use it in GitHub Desktop.
Segment Tree for range add and range sum
#pragma once
#include <vector>
#include <algorithm>
template<class T>
class segtree {
public:
struct node { T all; T part; };
explicit segtree(size_t N=1, T init=0)
: _init(init), _N(N), _tree(2*N-1,{init,init}) {}
void add(size_t i, size_t j, T data) {
_add(i,j,data,0,0,_N);
}
T sum(size_t i, size_t j) {
return _sum(i,j,0,0,_N);
}
private:
void _add(size_t i, size_t j, T data, size_t k, size_t l, size_t r) {
if(j <= l || r <= i) return;
if(i <= l && r <= j) {
_tree[k].all += data;
} else {
_tree[k].part += data*(std::min(j,r)-std::max(i,l));
_add(i,j,data,k*2+1,l,(l+r)/2);
_add(i,j,data,k*2+2,(l+r)/2,r);
}
}
T _sum(size_t i, size_t j, size_t k, size_t l, size_t r) {
if(j <= l || r <= i) return _init;
if(i <= l && r <= j) {
return _tree[k].part + _tree[k].all*(r-l);
} else {
const T vleft = _sum(i,j,k*2+1,l,(l+r)/2);
const T vright = _sum(i,j,k*2+2,(l+r)/2,r);
return _tree[k].all*(std::min(j,r)-std::max(i,l)) + vleft + vright;
}
}
private:
T _init;
size_t _N;
std::vector<node> _tree;
};
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment