SegmentTree (array)

SegmentTree.cc

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

using namespace std;

class SegmentTree {
private:
  // number of elements
  int _N;
  // segment tree nodes array
  vector<int> _arr;
  void build(vector<int> &arr, int curr, int left, int right);
  int query(int curr, int left, int right, int leftQuery, int rightQuery);
  void update(int curr, int left, int right, int pos, int new_val);
public:
  void build(vector<int> &arr);
  int query(int leftQuery, int rightQuery);
  void update(int pos, int newVal);
};

void SegmentTree::build(vector<int> &arr) {
  _N = arr.size();
  int height = static_cast<int>(ceil(log2(_N)));
  _arr = vector<int>(pow(2, height + 1) - 1);
  build(arr, 0, 0, _N - 1);
}

void SegmentTree::build(vector<int> &arr, int curr, int left, int right) {
  if (left == right) _arr[curr] = arr[left];
  else {
    int mid = (left + right) >> 1;
    build(arr, curr * 2 + 1, left, mid);
    build(arr, curr * 2 + 2, mid + 1, right);
    _arr[curr] = _arr[curr * 2 + 1] + _arr[curr * 2 + 2];
  }
}

int SegmentTree::query(int leftQuery, int rightQuery) {
  return query(0, 0, _N - 1, leftQuery, rightQuery);
}

int SegmentTree::query(int curr, int left, int right, int leftQuery, int rightQuery) {
  if (right < leftQuery || left > rightQuery) return 0;
  if (leftQuery <= left && rightQuery >= right) return _arr[curr];
  int mid = (left + right) >> 1;
  return query(curr * 2 + 1, left, mid, leftQuery, rightQuery) +
    query(curr * 2 + 2, mid + 1, right, leftQuery, rightQuery);
}

void SegmentTree::update(int curr, int left, int right, int pos, int newVal) {
  if (left == right) {
    _arr[curr] = newVal;
  } else {
    int mid = (left + right) / 2;
    if (pos <= mid)
      update(curr * 2 + 1, left, mid, pos, newVal);
    else
      update(curr * 2 + 2, mid + 1, right, pos, newVal);
    _arr[curr] = _arr[curr * 2 + 1] + _arr[curr * 2 + 2];
  }
}

void SegmentTree::update(int pos, int val) {
  update(0, 0, _N - 1, pos, val);
}