普通平衡树【模板】【treap】

学了一下treap的模板感觉还算容易理解。
其实就是平衡二叉树+堆
每个节点都有一个随机的权值。
根据堆的特性,只有右旋和左旋操作。
个人觉得关键代码就插入和删除部分涉及旋转的维护代码需要理解一下,其他的询问操作都很好写。

#pragma GCC optimize(2)
#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e5 + 5;
typedef long long ll;
const ll mod = 1e9 + 7;
int Case = 1;
int n, m;
struct Treap {
  int ch[maxn][2], dat[maxn], size[maxn], cnt[maxn], val[maxn];
  int root, tot;
  Treap() { tot = 0; }
  int newnode(int v) {
    val[++tot] = v;
    dat[tot] = rand();
    size[tot] = 1;
    cnt[tot] = 1;
    return tot;
  }
  void pushup(int rt) {
    size[rt] = cnt[rt] + size[ch[rt][0]] + size[ch[rt][1]];
  }
  void rotate(int &rt, int d) {
    int v = ch[rt][d];
    ch[rt][d] = ch[v][d ^ 1];
    ch[v][d ^ 1] = rt;
    rt = v;
    pushup(rt);
    pushup(ch[rt][d ^ 1]);
  }
  void insert(int &rt, int v) {
    if (!rt) {
      rt = newnode(v);
      return;
    }
    if (v == val[rt])
      cnt[rt]++;
    else {
      int d = v < val[rt] ? 0 : 1;
      insert(ch[rt][d], v);
      if (dat[rt] < dat[ch[rt][d]]) rotate(rt, d);
    }
    pushup(rt);
  }
  void del(int &rt, int v) {
    if (!rt) return;
    if (val[rt] == v) {
      if (cnt[rt] > 1) {
        cnt[rt]--;
        pushup(rt);
        return;
      }
      if (ch[rt][1] || ch[rt][0]) {
        if (!ch[rt][1] || dat[ch[rt][0]] > dat[ch[rt][1]])
          rotate(rt, 0), del(ch[rt][1], v);
        else
          rotate(rt, 1), del(ch[rt][0], v);
        pushup(rt);
      } else
        rt = 0;
      return;
    }
    v < val[rt] ? del(ch[rt][0], v) : del(ch[rt][1], v);
    pushup(rt);
  }
  int get_rank(int rt, int v) {
    if (!rt) return 0;
    if (v == val[rt])
      return size[ch[rt][0]] + 1;
    else if (v < val[rt])
      return get_rank(ch[rt][0], v);
    else
      return size[ch[rt][0]] + cnt[rt] + get_rank(ch[rt][1], v);
  }
  int get_val(int rt, int rank) {
    if (!rt) return -1;
    if (rank <= size[ch[rt][0]])
      return get_val(ch[rt][0], rank);
    else if (rank <= size[ch[rt][0]] + cnt[rt])
      return val[rt];
    else
      return get_val(ch[rt][1], rank - size[ch[rt][0]] - cnt[rt]);
  }
  int get_pre(int v) {
    int res = 0, rt = root;
    while (rt) {
      if (val[rt] < v)
        res = val[rt], rt = ch[rt][1];
      else
        rt = ch[rt][0];
    }
    return res;
  }
  int get_next(int v) {
    int res, rt = root;
    while (rt) {
      if (val[rt] > v)
        res = val[rt], rt = ch[rt][0];
      else
        rt = ch[rt][1];
    }
    return res;
  }
} treap;
void solve() {
  scanf("%d", &n);
  for (int i = 1; i <= n; i++) {
    int op;
    scanf("%d", &op);
    if (op == 1) {
      int val;
      scanf("%d", &val);
      treap.insert(treap.root, val);
    } else if (op == 2) {
      int val;
      scanf("%d", &val);
      treap.del(treap.root, val);
    } else if (op == 3) {
      int val;
      scanf("%d", &val);
      printf("%d\n", treap.get_rank(treap.root, val));
    } else if (op == 4) {
      int rank;
      scanf("%d", &rank);
      printf("%d\n", treap.get_val(treap.root, rank));
    } else if (op == 5) {
      int val;
      scanf("%d", &val);
      printf("%d\n", treap.get_pre(val));
    } else if (op == 6) {
      int val;
      scanf("%d", &val);
      printf("%d\n", treap.get_next(val));
    }
  }
  return;
}
int main() {
  // g++ -std=c++11 -o2 1.cpp -o f && ./f < in.txt
  // ios::sync_with_stdio(false);
#ifndef ONLINE_JUDGE
  // freopen("in.txt", "r", stdin);
  // freopen("out.txt","w",stdout);
#endif
  while (Case--) {
    solve();
  }
  return 0;
}

猜你喜欢

转载自blog.csdn.net/qq_39921637/article/details/95617177