数据结构-Treap

高级数据结构-Treap

前两天看了BYVoid大神写的Treap的论文，论文中几乎给出了所有的代码，不过最后一部提到了一个优化，就是对于重复的元素采用在同一个节点中记录一个weight来记录者个元素使用了多少次的写法，以及查找第K个数的功能的实现，我试着写了一个，表示树旋转之后需要重新设置size有一点点烦。因为找不到合适的题目去测试这段代码，主要是因为这里采用动态内存分配导致代码运行速度很慢总是TLE，我自己做了一下测试，不过目前只测试了插入功能和查找功能。删除功能，查找第K个数的功能，给定一个数给出排名的功能还没有测试过。目标是最终实现一个类似与STL中map的模板类，同时写出一个ACM竞赛中可用的treap模板（取消内存动态分配，采用数组静态存储）。

先把代码挂上来先，测试什么的慢慢来。

************************************************

已经证实一开始的代码是有错的，Treap的删除操作中，当找到节点时那个--cnt==0的判断条件需要改一下，因为整个过程涉及递归，在下一步递归时再次--，cnt就会变成负数。

这只能算个补丁，应该有更好的写法。

************************************************

# include <cstdio>
# include <cstring>
# include <cmath>
# include <cstdlib>
# include <climits>
# include <iostream>
# include <iomanip>
# include <set>
# include <map>
# include <vector>
# include <stack>
# include <queue>
# include <algorithm>
using namespace std;

const int debug = 0;
const int size  = 1000000 + 10;
typedef long long ll;

# ifndef ONLINE_JUDGE
struct DesktopIO{
	DesktopIO(){
		freopen("//home//unclesugar//in.txt","r",stdin);
		freopen("//home//unclesugar//out.txt","w",stdout);
	}
}DIO;
# endif

struct Treap_Node{
	int value;
	int fix,cnt,size;
	Treap_Node *left,*right;
	Treap_Node():cnt(0),size(0),left(NULL),right(NULL){}
	Treap_Node(int _value):value(_value),cnt(0),size(0),left(NULL),right(NULL){}
}*root = NULL;
inline void Treap_SetSize(Treap_Node *&P){
	if (P){
		P->size = P->cnt;
		if (P->left)	P->size += P->left->size;
		if (P->right)	P->size += P->right->size;
	}
}
inline int lsize(Treap_Node *&P){
	return P->left?P->left->size:0;
}
inline int rsize(Treap_Node *&P){
	return P->right?P->right->size:0;
}
void Treap_Left_Rotate(Treap_Node *&a){
	Treap_Node *b = a->right;
	a->right = b->left;
	b->left = a;
	a = b;
	Treap_SetSize(a->left);
	Treap_SetSize(a->right);
	Treap_SetSize(a);
}
void Treap_Right_Rotate(Treap_Node *&a){
	Treap_Node *b = a->left;
	a->left = b->right;
	b->right = a;
	a = b;
	Treap_SetSize(a->left);
	Treap_SetSize(a->right);
	Treap_SetSize(a);
}
void Treap_Insert(Treap_Node *&P,int value){
	if (!P){
		P = new Treap_Node;
		P->value = value;
		P->fix = rand();
	}
	if (value < P->value){
		Treap_Insert(P->left,value);
		if (P->left->fix < P->fix)
			Treap_Right_Rotate(P);
	}
	else if (P->value < value){
		Treap_Insert(P->right,value);
		if (P->right->fix < P->fix)
			Treap_Left_Rotate(P);
	}
	else {
		P->cnt++;
	}
	Treap_SetSize(P);
}
bool Treap_Delete(Treap_Node *&P,int value){
	bool ret = false;
	if (!P)	{
		ret = false;
	}
	else {
		if (value < P->value)
			Treap_Delete(P->left,value);
		else if (P->value < value)
			Treap_Delete(P->right,value);
		else {
			if (P->cnt==0||(--P->cnt)==0){
				if (!P->left||!P->right){
					Treap_Node *t = P;
					if (!P->right)
						P = P->left;
					else
						P = P->right;
					delete t;
					ret = true;
				}
				else if (P->left->fix < P->right->fix){
					Treap_Right_Rotate(P);
					ret = Treap_Delete(P->right,value);
				}
				else {
					Treap_Left_Rotate(P);
					ret = Treap_Delete(P->left,value);
				}
			}
		}
		Treap_SetSize(P);
	}
	return ret;
}

Treap_Node* Treap_Find(Treap_Node* root,int value){
	while (root){
		if (value < root->value)
			root = root->left;
		else if (root->value < value)
			root = root->right;
		else
			break;
	}
	return root;
}
Treap_Node* Treap_Findkth(Treap_Node *&P,int k){
	if (k <= lsize(P))
		return Treap_Findkth(P->left,k);
	else if (k > lsize(P)+P->cnt)
		return Treap_Findkth(P->right,k-(lsize(P)+P->cnt));
	else return P;
}
int Treap_Rank(Treap_Node *P,int value,int cur){
	if (value==P->value)
		return lsize(P) + cur + 1;
	else if (value < P->value)
		return Treap_Rank(P->left,value,cur);
	else return Treap_Rank(P->right,value,cur+lsize(P)+P->cnt);
}
void Treap_Clear(Treap_Node *&root){
	if (root->left)
		Treap_Clear(root->left);
	if (root->right)
		Treap_Clear(root->right);
	delete root;
	root = NULL;
}
int main()
{
	std::ios::sync_with_stdio(false);cin.tie(0);
	int i,j,k,tmp;
	int n,m;
	while (cin >> n >> m){
		while (n--){
			cin >> tmp;
			Treap_Insert(root,tmp);
		}
		while (m--){
			cin >> tmp;
			if (Treap_Find(root,tmp))
				cout << "YES\n";
			else 
				cout << "NO\n";
		}
	}
	return 0;
}