Skip to content

segtree.hpp

#include "noya/segtree.hpp"

View on GitHub

#ifndef NOYA_SEGTREE_HPP
#define NOYA_SEGTREE_HPP 1

#include <cassert>
#include <vector>

namespace noya {

/// @brief Segment tree for a monoid type.
/// Monoid must provide `using value_type`, `value_type unit()`, and
/// `value_type op(value_type, value_type)`.
template <class Monoid> struct segtree {
  using MX = Monoid;
  using S = typename MX::value_type;
  using value_type = S;

  int n = 0;
  int size = 1;
  std::vector<S> d;

  segtree() {}
  explicit segtree(int _n) { build(_n); }

  explicit segtree(const std::vector<S> &v) { build(v); }

  template <class F> segtree(int _n, F f) { build(_n, f); }

  void build(int _n) { build(_n, [](int) { return MX::unit(); }); }

  void build(const std::vector<S> &v) {
    build(int(v.size()), [&](int i) { return v[i]; });
  }

  template <class F> void build(int _n, F f) {
    n = _n;
    size = 1;
    while (size < n)
      size <<= 1;
    d.assign(size << 1, MX::unit());
    for (int i = 0; i < n; i++)
      d[size + i] = f(i);
    for (int i = size - 1; i >= 1; i--)
      update(i);
  }

  void set(int p, S x) {
    assert(0 <= p && p < n);
    p += size;
    d[p] = x;
    while (p >>= 1)
      update(p);
  }

  void multiply(int p, S x) {
    assert(0 <= p && p < n);
    p += size;
    d[p] = MX::op(d[p], x);
    while (p >>= 1)
      update(p);
  }

  S get(int p) const {
    assert(0 <= p && p < n);
    return d[p + size];
  }

  std::vector<S> get_all() const {
    return std::vector<S>(d.begin() + size, d.begin() + size + n);
  }

  S prod(int l, int r) const {
    assert(0 <= l && l <= r && r <= n);
    S sml = MX::unit(), smr = MX::unit();
    l += size;
    r += size;
    while (l < r) {
      if (l & 1)
        sml = MX::op(sml, d[l++]);
      if (r & 1)
        smr = MX::op(d[--r], smr);
      l >>= 1;
      r >>= 1;
    }
    return MX::op(sml, smr);
  }

  S all_prod() const { return d[1]; }

  template <class F> int max_right(int l, F f) const {
    assert(0 <= l && l <= n);
    assert(f(MX::unit()));
    if (l == n)
      return n;
    l += size;
    S sm = MX::unit();
    do {
      while ((l & 1) == 0)
        l >>= 1;
      if (!f(MX::op(sm, d[l]))) {
        while (l < size) {
          l <<= 1;
          if (f(MX::op(sm, d[l]))) {
            sm = MX::op(sm, d[l]);
            l++;
          }
        }
        return l - size;
      }
      sm = MX::op(sm, d[l++]);
    } while ((l & -l) != l);
    return n;
  }

  template <class F> int min_left(int r, F f) const {
    assert(0 <= r && r <= n);
    assert(f(MX::unit()));
    if (r == 0)
      return 0;
    r += size;
    S sm = MX::unit();
    do {
      --r;
      while (r > 1 && (r & 1))
        r >>= 1;
      if (!f(MX::op(d[r], sm))) {
        while (r < size) {
          r = (r << 1) | 1;
          if (f(MX::op(d[r], sm))) {
            sm = MX::op(d[r], sm);
            --r;
          }
        }
        return r + 1 - size;
      }
      sm = MX::op(d[r], sm);
    } while ((r & -r) != r);
    return 0;
  }

private:
  void update(int k) { d[k] = MX::op(d[k << 1], d[k << 1 | 1]); }
};

} // namespace noya

#endif // NOYA_SEGTREE_HPP
#include <cassert>
#include <vector>

namespace noya {

/// @brief Segment tree for a monoid type.
/// Monoid must provide `using value_type`, `value_type unit()`, and
/// `value_type op(value_type, value_type)`.
template <class Monoid> struct segtree {
  using MX = Monoid;
  using S = typename MX::value_type;
  using value_type = S;

  int n = 0;
  int size = 1;
  std::vector<S> d;

  segtree() {}
  explicit segtree(int _n) { build(_n); }

  explicit segtree(const std::vector<S> &v) { build(v); }

  template <class F> segtree(int _n, F f) { build(_n, f); }

  void build(int _n) { build(_n, [](int) { return MX::unit(); }); }

  void build(const std::vector<S> &v) {
    build(int(v.size()), [&](int i) { return v[i]; });
  }

  template <class F> void build(int _n, F f) {
    n = _n;
    size = 1;
    while (size < n)
      size <<= 1;
    d.assign(size << 1, MX::unit());
    for (int i = 0; i < n; i++)
      d[size + i] = f(i);
    for (int i = size - 1; i >= 1; i--)
      update(i);
  }

  void set(int p, S x) {
    assert(0 <= p && p < n);
    p += size;
    d[p] = x;
    while (p >>= 1)
      update(p);
  }

  void multiply(int p, S x) {
    assert(0 <= p && p < n);
    p += size;
    d[p] = MX::op(d[p], x);
    while (p >>= 1)
      update(p);
  }

  S get(int p) const {
    assert(0 <= p && p < n);
    return d[p + size];
  }

  std::vector<S> get_all() const {
    return std::vector<S>(d.begin() + size, d.begin() + size + n);
  }

  S prod(int l, int r) const {
    assert(0 <= l && l <= r && r <= n);
    S sml = MX::unit(), smr = MX::unit();
    l += size;
    r += size;
    while (l < r) {
      if (l & 1)
        sml = MX::op(sml, d[l++]);
      if (r & 1)
        smr = MX::op(d[--r], smr);
      l >>= 1;
      r >>= 1;
    }
    return MX::op(sml, smr);
  }

  S all_prod() const { return d[1]; }

  template <class F> int max_right(int l, F f) const {
    assert(0 <= l && l <= n);
    assert(f(MX::unit()));
    if (l == n)
      return n;
    l += size;
    S sm = MX::unit();
    do {
      while ((l & 1) == 0)
        l >>= 1;
      if (!f(MX::op(sm, d[l]))) {
        while (l < size) {
          l <<= 1;
          if (f(MX::op(sm, d[l]))) {
            sm = MX::op(sm, d[l]);
            l++;
          }
        }
        return l - size;
      }
      sm = MX::op(sm, d[l++]);
    } while ((l & -l) != l);
    return n;
  }

  template <class F> int min_left(int r, F f) const {
    assert(0 <= r && r <= n);
    assert(f(MX::unit()));
    if (r == 0)
      return 0;
    r += size;
    S sm = MX::unit();
    do {
      --r;
      while (r > 1 && (r & 1))
        r >>= 1;
      if (!f(MX::op(d[r], sm))) {
        while (r < size) {
          r = (r << 1) | 1;
          if (f(MX::op(d[r], sm))) {
            sm = MX::op(d[r], sm);
            --r;
          }
        }
        return r + 1 - size;
      }
      sm = MX::op(d[r], sm);
    } while ((r & -r) != r);
    return 0;
  }

private:
  void update(int k) { d[k] = MX::op(d[k << 1], d[k << 1 | 1]); }
};

} // namespace noya