【ACWing】2476. 树套树

题目地址：

https://www.acwing.com/problem/content/2478/

请你写出一种数据结构，来维护一个长度为$n$的数列，其中需要提供以下操作：
1 l r x，查询整数$x$在区间$[l,r]$内的排名（如果$x$出现多于$1$次，则求的是最小排名）。
2 l r k，查询区间$[l,r]$内排名为$k$的值。
3 pos x，将pos位置的数修改为$x$。
4 l r x，查询整数$x$在区间$[l,r]$内的前驱(前驱定义为小于$x$，且最大的数)。
5 l r x，查询整数$x$在区间$[l,r]$内的后继（后继定义为大于$x$，且最小的数）。
数列中的位置从左到右依次标号为$1\sim n$。
区间$[l,r]$表示从位置$l$到位置$r$之间（包括两端点）的所有数字。
区间内排名为$k$的值指区间内从小到大排在第$k$位的数值。（位次从$1$开始）

输入格式：
第一行包含两个整数$n,m$，表示数列长度以及操作次数。
第二行包含$n$个整数，表示数列。
接下来$m$行，每行包含一个操作指令，格式如题目所述。

输出格式：
对于所有操作1,2,4,5，每个操作输出一个查询结果，每个结果占一行。

数据范围：
$1\le n,m\le 5\times 10^4$,
$1\le l\le r\le n$,
$1\le pos\le n$,
$1\le k\le r-l+1$,
$0\le x\le 10^8$,
有序数列中的数字始终满足在$[0,10^8]$范围内，
数据保证所有操作一定合法，所有查询一定有解。

思路是树套树，外层是线段树维护区间信息，每个线段树节点内部存一个FHQ Treap，维护该区间的有序数列的信息。几个操作做法分别如下：
1、求$x$在$[l,r]$内的排名可以转化为求$[l,r]$中比$x$小的数的个数（个数加$1$即为排名）。用线段树可以将$[l,r]$分解为不相交的对数级别的区间，每个区间查询一下FHQ Treap里小于$x$的个数，将个数加总即可。
2、求$[l,r]$内排名为$k$的值，直接求比较困难，可以用二分。这里即求排名小于等于$k$的最大数，而排名可以用上面的1来做。由于数据范围是$[0,10^8]$，可以在这个区间二分。
3、将位置$p$的数修改为$x$，我们事先是已经将数列读入到一个数组里的，查询数组即知位置$p$的数是什么，接着我们在线段树里将包含位置$p$的区间都做一下更新即可，更新的方式是在FHQ Treap里将旧的数删掉，然后插入$x$；
4、求$x$在$[l,r]$里的前驱，即求$[l,r]$里小于$x$的最大数。线段树可以将$[l,r]$分解为不相交的对数级别的区间，每个区间查询一下FHQ Treap里小于$x$的最大数，然后在这些最大数里求最大的即可。
5、求$x$在$[l,r]$里的后继，与$4$类似。

有几个需要注意的地方：
1、FHQ Treap节点要开$n\log n+m$个，线段树有$\log n$层，每层的FHQ Treap节点要$n$个，所以线段树需要的总的FHQ Treap的节点数为$n\log n$；此外对于操作$3$，每次操作一下都会新开一个节点，所以总共就是$n\log n+m$个；
2、由于这里FHQ Treap维护的是有序数列，在建树的时候要将节点逐个插入，以保证有序性；
3、由于查前驱和后继的时候，可能出现某个区间的FHQ Treap并不存在某个数的前驱，方便起见，如果不存在前驱，则设前驱为负无穷；如果不存在后继，则设后继为正无穷。

代码如下：

#include <iostream>
using namespace std;

const int N = 1e5 + 10, INF = 1e9;
int n, m;
int a[N];

namespace FHQ {

struct Node {
#define lc(u) tr[u].l
#define rc(u) tr[u].r
#define val(u) tr[u].val
#define sz(u) tr[u].sz
  int l, r;
  int val, sz;
  int rnd;
} tr[N * 21];
int idx;
int get_node(int v) {
  tr[++idx] = {0, 0, v, 1, rand()};
  return idx;
}

void pushup(int u) { sz(u) = sz(lc(u)) + sz(rc(u)) + 1; }

void split(int u, int val, int &x, int &y) {
  if (!u) x = y = 0;
  else {
    if (val < val(u)) {
      y = u;
      split(lc(u), val, x, lc(u));
      pushup(y);
    } else {
      x = u;
      split(rc(u), val, rc(u), y);
      pushup(x);
    }
  }
}

int merge(int x, int y) {
  if (!x || !y) return x | y;
  else {
    if (tr[x].rnd > tr[y].rnd) {
      rc(x) = merge(rc(x), y);
      pushup(x);
      return x;
    } else {
      lc(y) = merge(x, lc(y));
      pushup(y);
      return y;
    }
  }
}

// 针对每个具体的FHQ Treap的函数要写在struct内部
struct FHQ {
  int root;

  void insert(int v) {
    int x, y, z;
    split(root, v - 1, x, z);
    y = get_node(v);
    root = merge(merge(x, y), z);
  }
  // 将Treap中值为w的节点删掉，再插入v
  void modify(int w, int v) {
    int x, y, z;
    split(root, w - 1, x, z);
    split(z, w, y, z);
    y = merge(lc(y), rc(y));
    root = merge(merge(x, y), z);
	insert(v);
  }

  // 求Treap中小于v的节点个数
  int get_less(int v) {
    int x, y;
    split(root, v - 1, x, y);
    int res = sz(x);
    root = merge(x, y);
    return res;
  }

  // 求Treap中v的前驱。如果不存在则返回负无穷
  int get_prev(int v) {
    int u = root, res = -INF;
    while (u) {
      if (v > val(u)) res = max(res, val(u)), u = rc(u);
      else u = lc(u);
    }
    return res;
  }

  // 求Treap中v的后继。如果不存在则返回正无穷
  int get_next(int v) {
    int u = root, res = INF;
    while (u) {
      if (v < val(u)) res = min(res, val(u)), u = lc(u);
      else u = rc(u);
    }
    return res;
  }
};

}  // namespace FHQ

namespace SegTree {

// 每个线段树节点里有一个FHQ Treap
struct Node {
  int l, r;
  FHQ::FHQ fhq;
} tr[N << 2];

void build(int u, int l, int r) {
  tr[u] = {l, r};
  // 将该区间的值插入该区间的那个FHQ Treap
  for (int i = l; i <= r; i++) tr[u].fhq.insert(a[i]);
  if (l == r) return;

  int mid = l + r >> 1;
  build(u << 1, l, mid), build(u << 1 | 1, mid + 1, r);
}

int get_less(int u, int l, int r, int x) {
  if (l <= tr[u].l && tr[u].r <= r) return tr[u].fhq.get_less(x);
  int mid = tr[u].l + tr[u].r >> 1;
  int res = 0;
  if (l <= mid) res += get_less(u << 1, l, r, x);
  if (r > mid) res += get_less(u << 1 | 1, l, r, x);
  return res;
}

int get_key_from_rank(int lq, int rq, int rk) {
  int l = 0, r = 1e8;
  while (l < r) {
    int mid = l + r + 1 >> 1;
    if (get_less(1, lq, rq, mid) + 1 <= rk) l = mid;
    else r = mid - 1;
  }

  return l;
}

void modify(int u, int pos, int x) {
  if (pos < tr[u].l || pos > tr[u].r) return;
  tr[u].fhq.modify(a[pos], x);
  int mid = tr[u].l + tr[u].r >> 1;
  if (pos <= mid) modify(u << 1, pos, x);
  else modify(u << 1 | 1, pos, x);
}

int get_prev(int u, int l, int r, int x) {
  if (l <= tr[u].l && tr[u].r <= r) return tr[u].fhq.get_prev(x);
  int mid = tr[u].l + tr[u].r >> 1;
  int res = -INF;
  if (l <= mid) res = max(res, get_prev(u << 1, l, r, x));
  if (r > mid) res = max(res, get_prev(u << 1 | 1, l, r, x));
  return res;
}

int get_next(int u, int l, int r, int x) {
  if (l <= tr[u].l && tr[u].r <= r) return tr[u].fhq.get_next(x);
  int mid = tr[u].l + tr[u].r >> 1;
  int res = INF;
  if (l <= mid) res = min(res, get_next(u << 1, l, r, x));
  if (r > mid) res = min(res, get_next(u << 1 | 1, l, r, x));
  return res;
}

}  // namespace SegTree

int main() {
  scanf("%d%d", &n, &m);
  for (int i = 1; i <= n; i++) scanf("%d", &a[i]);
  SegTree::build(1, 1, n);
  while (m--) {
    int op, l, r, pos, x;
    scanf("%d", &op);
    if (op == 1 || op == 2 || op == 4 || op == 5) scanf("%d%d%d", &l, &r, &x);
    else scanf("%d%d", &pos, &x);

    if (op == 1) printf("%d\n", SegTree::get_less(1, l, r, x) + 1);
    else if (op == 2) printf("%d\n", SegTree::get_key_from_rank(l, r, x));
    else if (op == 3) {
      SegTree::modify(1, pos, x);
      // 线段树更新完，也要把数组a更新一下，以便于下次的修改
      a[pos] = x;
    } else if (op == 4) printf("%d\n", SegTree::get_prev(1, l, r, x));
    else printf("%d\n", SegTree::get_next(1, l, r, x));
  }
}

操作1、3、4、5时间复杂度$O({\log }^{2}n)$，操作2时间$O(\log R{\log }^{2}n)$，$R$是数据范围，空间$O(n\log n)$。