树堆treap
BST最大的问题在于不能自平衡,最坏情况下会退化成链。在BST的基础上,引入一个权值属性rnd,要求二叉树所有节点的rnd满足堆的性质,通过该成员的随机性来保持平衡,这就是Treap。易见,Treap是一种随机性数据结构。
关于满足堆的性质,大根堆或小根堆都可以,这里选择大根堆。
结构定义
struct Node {
int key, rnd;
Node *son[2];
Node(int x):key(x), rnd(rand()) {
son[0] = son[1] = nullptr;
}
int cmp(int x) const {
if (x == key) return -1;
return x < key ? 0 : 1;
}
};
旋转节点
树的平衡是通过旋转来实现的,以下是旋转操作,其中d取0表示左旋,取1表示右旋。
void Rotate(Node *&o, int d) {
Node *k = o->son[d^1];
o->son[d^1] = k->son[d];
k->son[d] = o;
o = k;
}
添加节点
void Insert(Node *&o, int x) {
if (o) {
int d = o->cmp(x);
if (d == -1) return;
Insert(o->son[d], x);
if (o->son[d]->rnd > o->rnd)
Rotate(o, d^1);
} else o = new Node(x);
}
删除节点
void Delete(Node *&o, int x) {
if (!o) return;
int d = o->cmp(x);
if (d == -1) {
if (!o->son[0]) o = o->son[1];
else if (!o->son[1]) o = o->son[0];
else {
int d2 = o->son[0]->rnd > o->son[1]->rnd ? 1 : 0;
Rotate(o, d2);
Delete(o->son[d2], x);
}
} else Delete(o->son[d], x);
}
查找节点
bool Exist(Node *o, int x) {
if (!o) return false;
int d = o->cmp(x);
if (d == -1) return true;
return Exist(o->son[d], x);
}
二分搜索
void LowerBound(Node *o, int x, int &ans) {
if (!o) return;
if (o->key >= x) ans = min(ans, o->key);
int d = x < o->key ? 0 : 1;
LowerBound(o->son[d], x, ans);
}
void UpperBound(Node *o, int x, int &ans) {
if (!o) return;
if (o->key > x) ans = min(ans, o->key);
int d = x < o->key ? 0 : 1;
UpperBound(o->son[d], x, ans);
}
第K小与排名
除了上述增删查等平衡树基本操作之外,Treap还可以实现查找第K小元素,以及给定key,问其在集合中的排名情况。
为了实现这两项功能,需要额外再增加个siz成员,用于标识以当前节点为根的子树的节点总数。
struct Node {
int key, rnd, siz;
Node *son[2];
Node(int x):key(x),rnd(rand()),siz(1) {
son[0] = son[1] = nullptr;
}
int cmp(int x) {
if (x == key) return -1;
return x < key ? 0 : 1;
}
};
在旋转时需要对siz进行更新。
int Size(Node *o) {
return o ? o->siz : 0;
}
void Maintain(Node *o) {
o->siz = 1 + Size(o->son[0]) + Size(o->son[1]);
}
void Rotate(Node *&o, int d) {
Node *k = o->son[d^1];
o->son[d^1] = k->son[d];
k->son[d] = o;
Maintain(o); Maintain(k);
o = k;
}
添加、删除以及查找的代码都不变。
第K小
int Select(Node *o, int k) {
int t = 1 + Size(o->son[0]);
if (k < t) return Select(o->son[0], k);
if (k > t) return Select(o->son[1], k-t);
return o->key;
}
排名
int Rank(Node *o, int x) {
if (!o) return 0;
if (x <= o->key) return Rank(o->son[0], x);
return 1 + Size(o->son[0]) + Rank(o->son[1], x);
}
Rank()返回的是小于x的元素个数,1+Rank()即为实际排名。