splaytree.hpp

Go to the documentation of this file.

#pragma once
 
#include <cassert>
#include <vector>
 
#include "library/data_structure/standard_monoids.hpp"
 
// Modified from https://judge.yosupo.jp/submission/144167,
// https://judge.yosupo.jp/submission/136748,
// https://judge.yosupo.jp/submission/278235
 
namespace maomao90 {
using namespace std;
 
namespace internal::splaytree {
template <Monoid T, Lazy<T> L, bool store_reverse> struct Node {
  Node *l, *r;
  int sz;
  bool rev;
  T val, sum;
  L lz;
  Node(T val = T::id())
      : l(nullptr), r(nullptr), sz(1), rev(false), val(val), sum(val),
        lz(L::id()) {}
};
 
template <Monoid T, Lazy<T> L> struct Node<T, L, true> {
  Node *l, *r;
  int sz;
  bool rev;
  T val, sum, rev_sum;
  L lz;
  Node(T val = T::id())
      : l(nullptr), r(nullptr), sz(1), rev(false), val(val), sum(val),
        rev_sum(val), lz(L::id()) {}
};
} // namespace internal::splaytree

template <Monoid T, Lazy<T> L, bool store_reverse = false> struct SplayTree {
private:
  using splaytree = SplayTree<T, L, store_reverse>;
  using node = internal::splaytree::Node<T, L, store_reverse>;
  node *root;
 
  int size(node *v) { return !v ? 0 : v->sz; }
  void update(node *v) {
    v->sz = 1;
    v->sum = v->val;
    if constexpr (store_reverse) {
      v->rev_sum = v->val;
    }
    if (v->l) {
      v->sz += v->l->sz;
      v->sum = v->l->sum.merge(v->sum);
      if constexpr (store_reverse) {
        v->rev_sum = v->rev_sum.merge(v->l->rev_sum);
      }
    }
    if (v->r) {
      v->sz += v->r->sz;
      v->sum = v->sum.merge(v->r->sum);
      if constexpr (store_reverse) {
        v->rev_sum = v->r->rev_sum.merge(v->rev_sum);
      }
    }
  }
  void push_down(node *v) {
    if (!v) {
      return;
    }
    if (v->l) {
      propagate(v->l, v->lz);
    }
    if (v->r) {
      propagate(v->r, v->lz);
    }
    v->lz = L::id();
    if (v->rev) {
      if (v->l) {
        reverse(v->l);
      }
      if (v->r) {
        reverse(v->r);
      }
      v->rev = false;
    }
  }
  void propagate(node *v, L x) {
    v->lz = x.merge(v->lz);
    v->val = x.apply(v->val, 1);
    v->sum = x.apply(v->sum, v->sz);
    if constexpr (store_reverse) {
      v->rev_sum = x.apply(v->rev_sum, v->sz);
    }
  }
  void reverse(node *v) {
    swap(v->l, v->r);
    if constexpr (store_reverse) {
      swap(v->sum, v->rev_sum);
    }
    v->rev ^= 1;
  }
  node *rotate_right(node *v) {
    node *l = v->l;
    v->l = l->r;
    l->r = v;
    update(v);
    update(l);
    return l;
  }
  node *rotate_left(node *v) {
    node *r = v->r;
    v->r = r->l;
    r->l = v;
    update(v);
    update(r);
    return r;
  }
  node *splay_top_down(node *v, int k) {
    push_down(v);
    int szl = v->l ? v->l->sz : 0;
    if (k == szl) {
      return v;
    }
    if (k < szl) {
      push_down(v->l);
      int szll = size(v->l->l);
      if (k == szll) {
        v = rotate_right(v);
      } else if (k < szll) {
        v->l->l = splay_top_down(v->l->l, k);
        v = rotate_right(v);
        v = rotate_right(v);
      } else {
        v->l->r = splay_top_down(v->l->r, k - szll - 1);
        v->l = rotate_left(v->l);
        v = rotate_right(v);
      }
    } else {
      push_down(v->r);
      k -= szl + 1;
      int szrl = size(v->r->l);
      if (k == szrl) {
        v = rotate_left(v);
      } else if (k < szrl) {
        v->r->l = splay_top_down(v->r->l, k);
        v->r = rotate_right(v->r);
        v = rotate_left(v);
      } else {
        v->r->r = splay_top_down(v->r->r, k - szrl - 1);
        v = rotate_left(v);
        v = rotate_left(v);
      }
    }
    update(v);
    return v;
  }
  node *merge_inner(node *l, node *r) {
    if (!l || !r) {
      return !l ? r : l;
    }
    r = splay_top_down(r, 0);
    r->l = l;
    update(r);
    return r;
  }
  pair<node *, node *> split_inner(node *v, int k) {
    int n = size(v);
    if (k >= n) {
      return {v, nullptr};
    }
    v = splay_top_down(v, k);
    node *l = v->l;
    v->l = nullptr;
    update(v);
    return {l, v};
  }
  tuple<node *, node *, node *> split3_inner(node *v, int l, int r) {
    if (l == 0) {
      auto [b, c] = split_inner(v, r);
      return {nullptr, b, c};
    }
    v = splay_top_down(v, l - 1);
    auto [b, c] = split_inner(v->r, r - l);
    v->r = nullptr;
    update(v);
    return {v, b, c};
  }
  // Only can be used to merge if it was split using `split3_inner`. O(1)
  // compared to general `merge3_inner`.
  node *inv_split3_inner(node *a, node *b, node *c) {
    node *v = merge_inner(b, c);
    if (!a) {
      return v;
    }
    a->r = v;
    update(a);
    return a;
  }
  node *merge3_inner(node *a, node *b, node *c) {
    node *v = merge_inner(b, c);
    return merge_inner(a, v);
  }
  node *set_inner(node *v, int k, T x) {
    v = splay_top_down(v, k);
    v->val = x;
    update(v);
    return v;
  }
  node *get_inner(node *v, int k, T &x) {
    v = splay_top_down(v, k);
    x = v->val;
    return v;
  }
  node *update_inner(node *v, int l, int r, L x) {
    if (r == l) {
      return v;
    }
    auto [a, b, c] = split3_inner(v, l, r);
    propagate(b, x);
    return inv_split3_inner(a, b, c);
  }
  node *query_inner(node *v, int l, int r, T &res) {
    if (r == l) {
      return v;
    }
    auto [a, b, c] = split3_inner(v, l, r);
    res = b->sum;
    return inv_split3_inner(a, b, c);
  }
  template <typename P>
  node *max_right_inner(node *v, int l, P pred, int &res) {
    res = l;
    if (l == size(v)) {
      return v;
    }
    v = splay_top_down(v, l);
    if (!pred(v->val)) {
      return v;
    }
    res++;
    push_down(v);
    v->r = max_right_inner(v->r, v->val, pred, res);
    if (v->r) {
      v = rotate_left(v);
    }
    update(v);
    return v;
  }
  template <typename P>
  node *max_right_inner(node *v, T sum, P pred, int &res) {
    if (!v) {
      return v;
    }
    push_down(v);
    T lsum = sum;
    if (v->l) {
      lsum = lsum.merge(v->l->sum);
    }
    lsum = lsum.merge(v->val);
    if (pred(lsum)) {
      int szl = v->l ? v->l->sz : 0;
      res += szl + 1;
      v->r = max_right_inner(v->r, lsum, pred, res);
      if (v->r) {
        v = rotate_left(v);
      }
    } else {
      v->l = max_right_inner(v->l, sum, pred, res);
      if (v->l) {
        v = rotate_right(v);
      }
    }
    update(v);
    return v;
  }
  template <typename P> node *min_left_inner(node *v, int r, P pred, int &res) {
    res = r;
    if (r == 0) {
      return v;
    }
    v = splay_top_down(v, r - 1);
    if (!pred(v->val)) {
      return v;
    }
    res--;
    push_down(v);
    v->l = min_left_inner(v->l, v->val, pred, res);
    if (v->l) {
      v = rotate_right(v);
    }
    update(v);
    return v;
  }
  template <typename P> node *min_left_inner(node *v, T sum, P pred, int &res) {
    if (!v) {
      return v;
    }
    push_down(v);
    T rsum = sum;
    if (v->r) {
      rsum = v->r->sum.merge(rsum);
    }
    rsum = v->val.merge(rsum);
    if (pred(rsum)) {
      int szr = v->r ? v->r->sz : 0;
      res -= szr + 1;
      v->l = min_left_inner(v->l, rsum, pred, res);
      if (v->l) {
        v = rotate_right(v);
      }
    } else {
      v->r = min_left_inner(v->r, sum, pred, res);
      if (v->r) {
        v = rotate_left(v);
      }
    }
    update(v);
    return v;
  }
  node *reverse_inner(node *v, int l, int r) {
    if (r == l) {
      return v;
    }
    auto [a, b, c] = split3_inner(v, l, r);
    if (b) {
      reverse(b);
    }
    return inv_split3_inner(a, b, c);
  }
  node *insert_inner(node *v, int k, node *u) {
    if (k == size(v)) {
      u->l = v;
      update(u);
      return u;
    }
    if (k == 0) {
      u->r = v;
      update(u);
      return u;
    }
    v = splay_top_down(v, k);
    u->l = v->l;
    v->l = u;
    update(u);
    update(v);
    return v;
  }
  node *erase_inner(node *v, int k) {
    v = splay_top_down(v, k);
    return merge_inner(v->l, v->r);
  }
  node *build(const vector<T> &v, int l, int r) {
    int m = (l + r) >> 1;
    node *u = new node(v[m]);
    if (m > l) {
      u->l = build(v, l, m);
    }
    if (r > m + 1) {
      u->r = build(v, m + 1, r);
    }
    update(u);
    return u;
  }
  SplayTree(node *r) : root(r) {}
 
public:
  SplayTree() : root(nullptr) {}
  explicit SplayTree(int n) : SplayTree(vector<T>(n, T::id())) {}
  explicit SplayTree(const vector<T> &v) : root(nullptr) {
    if (!v.empty()) {
      root = build(v, 0, v.size());
    }
  }

  int size() { return size(root); }
  void set(int k, T x) {
    assert(0 <= k && k < size());
    root = set_inner(root, k, x);
  }

  T get(int k) {
    assert(0 <= k && k < size());
    T res = T::id();
    root = get_inner(root, k, res);
    return res;
  }

  void update(int l, int r, L x) {
    assert(0 <= l && l <= r && r <= size());
    root = update_inner(root, l, r, x);
  }

  T query(int l, int r) {
    assert(0 <= l && l <= r && r <= size());
    T res = T::id();
    root = query_inner(root, l, r, res);
    return res;
  }

  template <typename P> int max_right(int l, P pred) {
    assert(0 <= l && l <= size());
    int res = l;
    root = max_right_inner(root, l, pred, res);
    return res;
  }

  template <bool (*pred)(T)> int max_right(int l) {
    return max_right(l, [](T x) { return pred(x); });
  }

  template <typename P> int min_left(int r, P pred) {
    assert(0 <= r && r <= size());
    int res = r;
    root = min_left_inner(root, r, pred, res);
    return res;
  }

  template <bool (*pred)(T)> int min_left(int r) {
    return min_left(r, [](T x) { return pred(x); });
  }

  void reverse(int l, int r) {
    assert(0 <= l && l <= r && r <= size());
    root = reverse_inner(root, l, r);
  }

  void insert(int k, T x) {
    assert(0 <= k && k <= size());
    root = insert_inner(root, k, new node(x));
  }

  void erase(int k) {
    assert(0 <= k && k < size());
    root = erase_inner(root, k);
  }

  splaytree split(int k) {
    assert(0 <= k && k <= size());
    auto [a, b] = split_inner(root, k);
    root = a;
    return splaytree(b);
  }

  pair<splaytree, splaytree> split(int l, int r) {
    assert(0 <= l && l <= r && r <= size());
    auto [a, b, c] = split3_inner(root, l, r);
    root = a;
    return {splaytree(b), splaytree(c)};
  }

  void merge(splaytree &o) {
    root = merge_inner(root, o.root);
    o.root = nullptr;
  }

  void merge(splaytree &b, splaytree &c) {
    root = merge3_inner(root, b.root, c.root);
    b.root = c.root = nullptr;
  }
};
} // namespace maomao90