Treap学习

今天做PAT1057, 感觉平衡二叉树可以搞过， 学习了下treap
下面以PAT1057为例 （参考liuruja的书）,代码记在这里备忘

Treap

Treap可以动态维护一个有序表，支持在O(logn)的时间内完成插入一个元素，删除一个元素和查找第k大元素
Treap是一种平衡化的二叉搜索树，正如它的名字一样， 结合了tree + heap = treap
Treap是一种拥有键值和优先级两种权值的树，正是靠优先级来平衡的二叉树， 可以证明，如果优先级是随机的，treap的期望深度就是O(logn)
结点满足堆性质： 任意节点的优先级大于等于它子节点的优先级
treap节点的优先级是随机确定的， 可以用C/C++中的rand()函数来产生随机优先级

节点定义

struct Node {
	Node *ch[2]; //左右儿子, 这样用下标0和1可以把左旋右旋两种旋转统一起来简化代码
	int rank;   //优先级
	int value;  //键值
	int cnt;   //保存当前值的个数 
	int size; //本子树结点的个数

	Node():rank(0), value(0), cnt(0), size(0) {
		ch[0] = ch[1] = NULL;
	}

	bool operator < (const Node& other)  const {
		return rank < other.rank;
	}

	int cmp(int x) const {   //和当前节点的值比较， 相等返回-1， 返回0表示去左子树找，1表示去右子树找
		if (x == value) {
			return -1;
		}
		return x < value? 0 : 1;
	}

	void maintain() {  //保持结点的个数正确
		size = cnt;
		if (ch[0] != NULL) size += ch[0]->size;
		if (ch[1] != NULL) size += ch[1]->size;
	}
};

旋转

//画一下图就明白了, 这里的d只能是0或者1
void rotate(Node* &o, int d) { // d表示旋转方向， 0 左旋  1右旋  
	Node* k = o->ch[d ^ 1];   //左旋保存右子树， 右旋保存左子树
	o->ch[d ^ 1] = k->ch[d];
	k->ch[d] = o;
	o->maintain(); //这里o和k的顺序不能换， 因为o已经成为k的子树了
	k->maintain();
	o = k;
}

插入

void insert(Node* &o, int x) {
	if (o == NULL) {
		o = new Node();
		o->value = x;
		o->cnt = 1;
		o->rank = rand();
	} else {
		int d = o->cmp(x);
		if (d == -1) {
			o->cnt++;
		} else {
			insert(o->ch[d], x);
			if (o->rank < o->ch[d]->rank) { //如果子树的优先级大于当前节点的优先级，就要旋转，左边大于就右旋， 右边大于就左旋
				rotate(o, d ^ 1);
			}
		}
	}
	o->maintain(); //别忘了调整size
}

删除

void remove(Node* &o, int x) {
	int d = o->cmp(x);
	if (d == -1) {
		Node* u = o;
		if (o->cnt > 1) {
			o->cnt--;
		} else if (o->ch[0] != NULL && o->ch[1] != NULL) { 有两颗子树，把当前节点下移直到有一个结点或者到达叶子结点
			int d2 = (o->ch[0]->rank > o->ch[1]->rank ? 1 : 0); //这里把优先级高的旋转成根
			rotate(o, d2);
			remove(o->ch[d2], x);
		} else { //只有一颗子树， 直接用子树替换当前结点就可以
			if (o->ch[0] == NULL) o = o->ch[1];
			else o = o->ch[0];
			delete u;
		}
	} else { 
		remove(o->ch[d], x);
	}
	if (o != NULL) o->maintain(); //别忘了调整size
}

查找

//查找就是普通的排序二叉树查找
int find(Node* o, int x) {
	while(o != NULL) {
		int d = o->cmp(x);
		if (d == -1) return 1;
		else o = o->ch[d];
	}
	return 0;
}

查找第k大的数

//很好理解吧
int kth(Node* o, int k) {
	if (o == NULL || k <= 0 || k > o->size) return 0;
	int s = o->ch[0] == NULL ? 0 : o->ch[0]->size;
	if (k > s && k <= s + o->cnt) {
		return o->value;
	} else if (k <= s) {
		return kth(o->ch[0], k);
	} else return kth(o->ch[1], k - s - o->cnt);
}

删除整颗树

void removetree(Node* &root) {
	if (root->ch[0] != NULL) removetree(root->ch[0]);
	if (root->ch[1] != NULL) removetree(root->ch[1]);
	delete root;
	root = NULL;
}