求助,严重超时的数组AVL
查看原帖
求助,严重超时的数组AVL
576807
URbit楼主2022/10/29 09:40

如题,代码如下 由于在网上找不到用数组写的AVL所以自己写了一下,感觉思路应该没错,不知道是哪里导致的超时

#include <iostream>

using std::cin;
using std::cout;
using std::endl;

const int MAX_SIZE = 1e7;

class AVLTree
{
public:
	int key[MAX_SIZE];	/*关键字*/
	int data[MAX_SIZE];	/*数据*/
	int bf[MAX_SIZE];	/*平衡因子*/
	int height[MAX_SIZE];/*高度*/
	int cnt[MAX_SIZE];	/*重复元素个数*/
	int size[MAX_SIZE];	/*元素数*/

	int pa[MAX_SIZE];	//父亲结点
	int leftSon[MAX_SIZE];	//左儿子
	int rightSon[MAX_SIZE];	//右儿子
	
	int root = 0;	//根结点
	int sp = 0;	//空间指针

	void updata(int x);	/*更新信息*/
	void leftRotation(int x);	/*左旋*/
	void rightRotation(int x);	/*右旋*/
	int find(int value); /*查找*/
	void maintain(int x);	/*维持平衡*/
	void insert(int value);	/*插入*/
	void del(int value);	/*删除*/
	int kth_min(int k);	/*找第k小*/
	int kth_max(int k);	/*找第k大*/
	int findPrecursor(int value);	/*找前驱*/
	int findSuccessor(int value);	/*找后继*/
	int findTh(int value);	/*找数的排名*/
};

void AVLTree::updata(int x)
{
	/*
	* 更新结点信息
	*/

	if (x == 0) return;

	int p = pa[x];
	height[p] = height[leftSon[p]] > rightSon[p] ? height[leftSon[p]] + 1 : height[rightSon[p]] + 1;
	height[x] = height[leftSon[x]] > rightSon[x] ? height[leftSon[x]] + 1 : height[rightSon[x]] + 1;
	bf[p] = height[leftSon[p]] - height[rightSon[p]];
	bf[x] = height[leftSon[x]] - height[rightSon[x]];
	size[x] = cnt[x];
	if (leftSon[x]) size[x] += size[leftSon[x]];
	if (rightSon[x]) size[x] += size[rightSon[x]];

	root = rightSon[0];
}

void AVLTree::leftRotation(int x)
{
	/*
	* 左旋
	* 旋下去
	*/

	if (x == 0 || pa[x] == 0) return;

	int p = pa[x], r = rightSon[x];
	rightSon[x] = leftSon[r];
	leftSon[r] = x;
	pa[r] = p;
	pa[x] = r;
	if (rightSon[p] == x)
		rightSon[p] = r;
	else
		leftSon[p] = r;
	
	updata(x);
}

void AVLTree::rightRotation(int x)
{
	/*
	* 右旋
	* 旋下去
	*/

	if (x == 0 || pa[x] == 0) return;

	int p = pa[x], l = leftSon[x];
	leftSon[x] = rightSon[l];
	rightSon[l] = x;
	pa[l] = p;
	pa[x] = l;
	if (rightSon[p] == x)
		rightSon[p] = l;
	else
		leftSon[p] = l;

	updata(x);
}

int AVLTree::find(int value)
{
	/*
	* 查询结点
	*/

	int p = root;
	while (p)
	{
		if (value > key[p])
			p = rightSon[p];
		else if (value < key[p])
			p = leftSon[p];
		else
			return p;
	}
	return -1;
}

void AVLTree::maintain(int x)
{
	/*
	* 维持树高
	*/

	if (x == 0) return;

	int l = leftSon[x];
	int r = rightSon[x];

	int leftHeight = height[l];
	int rightHeight = height[r];

	if (leftHeight - rightHeight < -1 || leftHeight - rightHeight > 1)
	{
		if (leftHeight > rightHeight)
		{
			int llHeight = l == 0 ? 0 : height[leftSon[l]];
			int rlHeight = l == 0 ? 0 : height[rightSon[l]];
			if (llHeight >= rlHeight)
			{
				/*右旋一次*/
				rightRotation(x);
			}
			else
			{
				/*左右双旋*/
				leftRotation(l);
				rightRotation(x);
			}
		}
		else
		{
			int lrHeight = r == 0 ? 0 : height[leftSon[r]];
			int rrHeight = r == 0 ? 0 : height[rightSon[r]];
			if (lrHeight >= rrHeight)
			{
				/*左旋一次*/
				leftRotation(x);
			}
			else
			{
				/*右左双旋*/
				rightRotation(l);
				leftRotation(x);
			}
		}
	}
	updata(x);
}

void AVLTree::insert(int value)
{
	/*
	* 插入结点
	*/

	if (root == 0)
	{
		/*若此时树大小为0*/
		root = ++sp;
		pa[sp] = 0;
		rightSon[0] = sp;
		key[sp] = value;
		height[sp] = 0;
		cnt[sp]++;
		size[sp]++;
		return;
	}
	int p = root;
	int pp = 0;
	while (p != 0)
	{
		pp = p;
		if (key[p] > value)
			p = leftSon[p];
		else if (key[p] < value)
			p = rightSon[p];
		else
		{
			/*若找到相应的结点*/
			cnt[p]++;
			while (p)
			{
				size[p]++;
				p = pa[p];
			}
			return;
		}
	}

	/*更新结点信息*/
	++sp;
	pa[sp] = pp;
	key[sp] = value;
	height[sp] = 0;
	cnt[sp]++;
	size[sp]++;
	if (value < key[pp])
		leftSon[pp] = sp;
	else
		rightSon[pp] = sp;

	/*维持结点平衡*/
	p = sp;
	while (p != 0)
	{
		pp = pa[p];
		maintain(p);
		p = pp;
	}
}

void AVLTree::del(int value)
{
	/*
	* 删除结点(真删)
	*/
	int x = find(value);
	if (x == -1) return;

	/*找到了*/
	cnt[x]--;
	if (cnt[x] == 0)
	{
		/*用于替换的结点*/
		int replace = 0;

		/*将左右儿子均为0的情况统一到下面这两项*/
		if (leftSon[x] == 0)
		{
			/*左儿子为0*/
			replace = rightSon[x];
		}
		else if (rightSon[x] = 0)
		{
			/*右儿子为0*/
			replace = leftSon[x];
		}
		else
		{
			/*双儿子均不为0*/
			if (height[leftSon[x]] >= height[rightSon[x]])
			{
				/*左树高*/
				replace = leftSon[x];
				if (rightSon[replace])
				{
					while (rightSon[replace])
					{
						replace = rightSon[replace];
					}
					int p = pa[replace];
					rightSon[p] = leftSon[replace];
					leftSon[replace] = leftSon[x];
					rightSon[replace] = rightSon[x];
				}
				else
				{
					rightSon[replace] = rightSon[x];
				}
			}
			else
			{
				/*右树高*/
				replace = rightSon[x];
				if (leftSon[replace])
				{
					while (leftSon[replace])
					{
						replace = leftSon[replace];
					}
					int p = pa[replace];
					leftSon[p] = rightSon[replace];
					leftSon[replace] = leftSon[x];
					rightSon[replace] = rightSon[x];
				}
				else
				{
					leftSon[replace] = leftSon[x];
				}
			}
		}

		/*调整与父亲结点的关系*/
		int p = pa[replace];
		pa[replace] = pa[x];
		if (leftSon[pa[x]] = x)
		{
			leftSon[pa[x]] = replace;
			pa[x] = -1;
		}
		else
		{
			rightSon[pa[x]] = replace;
			pa[x] = -1;
		}

		/*维持结点*/
		while (pa[p])
		{
			int pp = pa[p];
			maintain(p);
			p = pp;
		}
	}
	else
	{
		int p = x;
		while (p != 0)
		{
			size[p]--;
			p = pa[x];
		}
	}
}

int AVLTree::kth_min(int k)
{
	/*
	* 找第k小
	*/

	int p = root;
	int lsize;
	if (p == 0 || k > size[p]) return 0;

	while (p)
	{
		lsize = size[leftSon[p]];
		if (k <= lsize) p = leftSon[p];
		else if (k <= lsize + cnt[p])
		{
			return key[p];
		}
		else
		{
			k -= (lsize + cnt[p]);
			p = rightSon[p];
		}
	}
}

int AVLTree::kth_max(int k)
{
	/*
	* 找第k大
	*/

	int p = root;
	int rsize;
	if (p == 0 || k > size[p]) return 0;

	while (p)
	{
		rsize = size[rightSon[p]];
		if (k <= rsize) p = rightSon[p];
		else if (k <= rsize + cnt[p])
		{
			return key[p];
		}
		else
		{
			k -= (rsize + cnt[p]);
			p = leftSon[p];
		}
	}
}

int AVLTree::findPrecursor(int value)
{
	/*
	* 找前驱
	*/

	int p = find(value);
	if (p == -1)
	{
		int p = root;
		int pp = p;
		while (p)
		{
			pp = p;
			if (value > key[p])
				p = rightSon[p];
			else if (value < key[p])
				p = leftSon[p];
		}
		if (key[pp] < value)
			return key[pp];
		else
			p = pp;
	}
	if (leftSon[p] != 0)
	{
		while (rightSon[p] != 0)
			p = rightSon[p];
	}
	else
	{
		while (p != rightSon[pa[p]])
			p = pa[p];
		p = pa[p];
		if (leftSon[p] != 0)
			while (rightSon[p] != 0)
				p = rightSon[p];
	}

	return key[p];
}

int AVLTree::findSuccessor(int value)
{
	/*
	* 找后继
	*/

	int p = find(value);
	if (p == -1)
	{
		int p = root;
		int pp = p;
		while (p)
		{
			pp = p;
			if (value > key[p])
				p = rightSon[p];
			else if (value < key[p])
				p = leftSon[p];
		}
		if (key[pp] > value)
			return key[pp];
		else
			p = pp;
	}
	if (rightSon[p] != 0)
	{
		while (leftSon[p] != 0)
			p = leftSon[p];
	}
	else
	{
		while (p != leftSon[pa[p]])
			p = pa[p];
		p = pa[p];
		if (rightSon[p] != 0)
			while (leftSon[p] != 0)
				p = leftSon[p];
	}
	return key[p];
}

int AVLTree::findTh(int value)
{
	/*
	* 找数的排名
	*/

	int p = find(value);
	int pp = p;
	int th = 0;
	while (rightSon[pa[pp]] = pp)
	{
		if (leftSon[pp])
			th += size[leftSon[pp]];
		th++;
	}
	if (leftSon[p])
		th += size[leftSon[p]];
	return th + 1;
}

AVLTree tree;

int main()
{
	int n;
	cin >> n;
	for (int i = 0; i < n; i++)
	{
		int opt, x;
		cin >> opt >> x;
		switch (opt)
		{
		case 1:
			tree.insert(x);
			break;
		case 2:
			tree.del(x);
			break;
		case 3:
			cout << tree.findTh(x) << "\n";
			break;
		case 4:
			cout << tree.kth_min(x) << "\n";
			break;
		case 5:
			cout << tree.findPrecursor(x) << "\n";
			break;
		case 6:
			cout << tree.findSuccessor(x) << "\n";
			break;
		}
	}

	return 0;
}
2022/10/29 09:40
加载中...