如题,代码如下 由于在网上找不到用数组写的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;
}