实现功能:
- 插入数值
- 删除数值
- 查询某排名的数字
- 查询某数值的排名
- 查询前驱后继
const int N = 100010, INF = 1e8;int n;
struct Node
{int l, r; // 左右子结点编号int key, val; // key:结点本身的值 val:为了使二叉树平衡的随机数int cnt, size; // cnt:当前结点的数出现了多少次 size:以当前结点为根的子树中有多少结点
}tr[N];int root, idx; // root:根结点编号 idx:计数器,记录结点编号用到哪了void pushup(int p) // 更新p
{tr[p].size = tr[tr[p].l].size + tr[tr[p].r].size + tr[p].cnt;
}int getnode(int key) // 创建值为key的新结点
{tr[ ++ idx].key = key;tr[idx].cnt = 1;tr[idx].size = 1;tr[idx].val = rand();return idx;
}void zig(int &p) // 右旋 (以p为根)
{int q = tr[p].l;tr[p].l = tr[q].r;tr[q].r = p;p = q;pushup(tr[p].r), pushup(p);
}void zag(int &p) // 左旋 (以p为根)
{int q = tr[p].r;tr[p].r = tr[q].l;tr[q].l = p;p = q;pushup(tr[p].l), pushup(p);
}void build() // 创建treap
{getnode(-INF), getnode(INF);root = 1, tr[1].r = 2;pushup(root);if (tr[1].val < tr[2].val) zag(root);
}void insert(int &p, int key) // 当前在p结点,要创建值为key的结点
{if (!p) p = getnode(key); // p不存在else if (tr[p].key == key) tr[p].cnt ++ ; // 原本就有key的结点else if (tr[p].key > key) // 当前结点的key比新结点的key大{insert(tr[p].l, key);if (tr[tr[p].l].val > tr[p].val) zig(p);}else // 当前结点的key比新结点的key小{insert(tr[p].r, key);if (tr[tr[p].r].val > tr[p].val) zag(p);}pushup(p); // 更新当前结点
}void remove(int &p, int key) // 当前在p结点,要删去值为key的结点
{if (!p) return;else if (tr[p].key == key) // 找到要删去的结点{if (tr[p].cnt > 1) tr[p].cnt -- ; // 想删去的key不止一个else if (tr[p].l || tr[p].r) // 左右子结点至少存在一个{if (!tr[p].r || tr[tr[p].l].val > tr[tr[p].r].val){zig(p);remove(tr[p].r, key);}else if (!tr[p].l || tr[tr[p].l].val < tr[tr[p].r].val){zag(p);remove(tr[p].l, key);}}else p = 0; // 叶子结点}else if (tr[p].key > key) remove(tr[p].l, key); // 当前结点值小于要删去的keyelse if (tr[p].key < key) remove(tr[p].r, key); // 当前结点值大于要删去的keypushup(p);
}int getrank_bykey(int p, int key) // 当前结点为p,给定key找rank
{if (!p) return 0;else if (tr[p].key == key) return tr[tr[p].l].size + 1; // 当前结点即为要找的结点else if (tr[p].key > key) return getrank_bykey(tr[p].l, key); // 要找的结点在当前左子树中else if (tr[p].key < key) return tr[tr[p].l].size + tr[p].cnt + getrank_bykey(tr[p].r, key); // 要找的结点在当前右子树中
}int getkey_byrank(int p, int rank) // 当前结点为p,给定rank找key
{if (!p) return INF;else if (tr[tr[p].l].size >= rank) return getkey_byrank(tr[p].l, rank); // 要找的结点在当前左子树else if (tr[tr[p].l].size + tr[p].cnt >= rank) return tr[p].key; // 要找的结点即为当前结点else return getkey_byrank(tr[p].r, rank - tr[tr[p].l].size - tr[p].cnt); // 要找的结点在当前右子树
}int getprev(int p, int key) // 当前结点为p,找小于key的最大值
{if (!p) return -INF;else if (tr[p].key >= key) return getprev(tr[p].l, key); // 小于key的结点在当前结点的左子树else if (tr[p].key < key) return max(tr[p].key, getprev(tr[p].r, key)); // 小于key的结点在当前结点or右子树
}int getnext(int p, int key) // 当前结点为p,找大于key的最小值
{if (!p) return INF;else if (tr[p].key <= key) return getnext(tr[p].r, key); // 大于key的结点在当前结点的右子树else if (tr[p].key > key) return min(tr[p].key, getnext(tr[p].l, key)); // 大于key的结点在当前结点or右子树
}