Segment Tree Sum
unknown
c_cpp
4 years ago
3.2 kB
18
Indexable
class segmentTree {
public:
long long* tree;
long long* lazyTree;
int N;
int size;
void printTree() {
for(int i = 0; i < size; i++) {
cout << tree[i] << ' ';
}
cout << '\n';
}
void update(int qleft, int qright, long long delta) {
update(qleft, qright, delta, 0, N - 1, 0);
}
void update(int qleft, int qright, long long delta, int left, int right, int index) {
// If changes at this node haven't been updated
if(lazyTree[index] != 0) {
long long change = lazyTree[index];
tree[index] += (right - left + 1)*lazyTree[index];
lazyTree[index] = 0;
// If not at a leaf, propagate to children
if(left != right) {
lazyTree[2*index + 1] += change;
lazyTree[2*index + 2] += change;
}
}
// Complete overlap
if(left >= qleft && right <= qright) {
tree[index] += (right - left + 1)*delta;
// If not a leaf
if(left != right) {
lazyTree[2*index + 1] += delta;
lazyTree[2*index + 2] += delta;
}
}
// Partial overlap
else if(!(right < qleft || left > qright)) {
int mid = (left + right)/2;
update(qleft, qright, delta, left, mid, 2*index + 1);
update(qleft, qright, delta, mid + 1, right, 2*index + 2);
tree[index] = tree[2*index + 1] + tree[2*index + 2];
}
}
long long query(int qleft, int qright) {
return query(qleft, qright, 0, N - 1, 0);
}
long long query(int qleft, int qright, int left, int right, int index) {
// If changes at this node haven't been updated
if(lazyTree[index] != 0) {
long long change = lazyTree[index];
tree[index] += (right - left + 1)*change;
lazyTree[index] = 0;
// If not at a leaf, propagate to children
if(left != right) {
lazyTree[2*index + 1] += change;
lazyTree[2*index + 2] += change;
}
}
// Complete overlap
if(left >= qleft && right <= qright) {
return tree[index];
}
// No overlap
if(right < qleft || left > qright) {
return 0;
}
// Partial
int mid = (left + right)/2;
return query(qleft, qright, left, mid, 2*index + 1) +
query(qleft, qright, mid + 1, right, 2*index + 2);
}
void createTree(vector<long long>& arr, int left, int right, int index) {
// Leaf
if(left == right) {
tree[index] = arr[left];
}
// Non-leaf
else {
int mid = (left + right)/2;
createTree(arr, left, mid, 2*index + 1);
createTree(arr, mid + 1, right, 2*index + 2);
tree[index] = tree[2*index + 1] + tree[2*index + 2];
}
}
segmentTree(vector<ll>& arr) {
N = arr.size();
// Find next highest 2^x - 1
int start = 2;
while((start - 1) < 2*N) start *= 2;
size = start - 1;
// Create space, initialize to 0 for range sum query
tree = new long long[size];
for(int i = 0; i < size; i++) {
tree[i] = 0;
}
// Same for lazy tree, init to 0
lazyTree = new long long[size];
memset(lazyTree, 0, size*sizeof(long long));
// Create tree
createTree(arr, 0, N - 1, 0);
}
};
Editor is loading...