lazy_segtree.cc

#include <iostream>
#include <vector>
#include <algorithm>
#include <string>

template<typename T, typename BinOp>
class lazy_segtree {
  struct node{
    int l, r;
    T val;
    T delayed;
  };

  int n;
  std::vector<node> dat;
  T init_val;
  BinOp op;

  public:
  lazy_segtree(): n(0) {}
  lazy_segtree(int n, T init_val, const BinOp &op): op(op), init_val(init_val) {
    int p = 1;
    while (p < n) p <<= 1;
    this->n = n = p;
    dat = std::vector<node>(p*2);
    for (int i=0; i<n; i++) {
      dat[n + i].l = i;
      dat[n + i].r = i + 1;
      dat[n + i].val = init_val;
      dat[n + i].delayed = init_val;
    }
    for (int i=n-1; i>0; i--) {
      dat[i].l = dat[2*i].l;
      dat[i].r = dat[2*i+1].r;
      dat[i].val = init_val;
      dat[i].delayed = init_val;
    }
  }

  void update_interval(int i, int j, T val) {
    update_interval(i, j, 1, val);
  }

  T query(int x, int y) {
    return query(x, y, 1);
  }

  private:
  T pow(const T &base, int n) {
    T acc = init_val;
    for (int i = 0; i < n; i++)
      acc = op(acc, base);
    return acc;
  }

  void update_interval(int x, int y, int c, const T &val) {
    node &cur = dat[c];

    if (x == cur.l && y == cur.r) {
      cur.delayed = op(cur.delayed, val);
      return;
    }

    int mid = dat[2*c].r;

    cur.val = op(cur.val, pow(val, y - x));
    if (x >= cur.l && y <= dat[2*c].r) {
      update_interval(x, y, 2*c, val);
      return;
    }

    if (y <= cur.r && x >= dat[2*c+1].l) {
      update_interval(x, y, 2*c+1, val);
      return;
    }

    update_interval(x, mid, 2*c, val);
    update_interval(mid, y, 2*c+1, val);
  }

  T query(int x, int y, int c) {
    node &cur = dat[c];

    cur.val = op(cur.val, pow(cur.delayed, cur.r - cur.l));
    if (c < this->n) {
      dat[2*c].delayed = op(dat[2*c].delayed, cur.delayed);
      dat[2*c+1].delayed = op(dat[2*c+1].delayed, cur.delayed);
    }
    cur.delayed = init_val;

    if (x == cur.l && y == cur.r) {
      return cur.val;
    }

    int mid = dat[2*c].r;

    if (x >= cur.l && y <= mid) {
      return query(x, y, 2*c);
    }

    if (y <= cur.r && x >= mid) {
      return query(x, y, 2*c+1);
    }

    return op(
        query(x, mid, 2*c),
        query(mid, y, 2*c+1));
  }
};

template<typename T, typename BinOp>
lazy_segtree<T, BinOp> build_lazy_segtree(int N, T init_val, BinOp op) {
  return lazy_segtree<T, BinOp>(N, init_val, op);
}


int main(int argc, char const* argv[])
{
  using namespace std;
  auto &&tree = build_lazy_segtree(100, 0, [](int x, int y){ return x + y; });

  tree.update_interval(40, 50, 1);
  cout << tree.query(0, 100) << endl;
  tree.update_interval(45, 60, 2);
  cout << tree.query(43, 54) << endl;
  
  return 0;
}