AVL树实现动态集合寻找第k小数

如果是在静态数组中寻找第k小数的话，可以利用快排思想，在快排代码基础上更改一些，便可完成，时间复杂度为O(n)，具体实现不展示。

但是在一个动态集合中，有着频繁的更新和删除操作的话，其实静态数组来实现并不合适，因为插入和删除操作的并不方便，以下是利用AVL树来实现。

具体操作为：在AVL树的基础下，在树中的每个节点中添加一个域——以该节点为根节点的子树的节点数，设为subTreeSize，其它域同AVL树。对于插入操作（函数addNewNode），在插入新节点节点的过程中每遍历到一个节点，更新其子树大小（即插入后子树大小加一）；对于删除操作（函数deleteOldNode），如正操的二叉查找树的插入操作，只是在删除递归过程中更新节点的子树大小（即插入后子树大小减一），插入和删除操作时间复杂度均为O(log n)（即树高）。然后找到第k小的操作(findKth)，若往左边遍历k的值不发生变化，且节点的值的排序大小为其左子树的节点数+1，而往右遍历时，由于左边子树及当前根节点的值都小于要找的第k小值，此时k-当前节点的排序大小（即在右子树里的相对大小），时间复杂度也为O(log n)，具体实现代码如下：

数据结构的定义

typedef struct AVLNode* Position;
typedef Position AVLTree; /* AVL树类型 */
struct AVLNode {
	int value;
	AVLTree left;
	AVLTree right;
	int height;       //树高
	int subTreeSize;//以当前节点为根节点的子树大小
};

所用到的函数

Position findMin(AVLTree Tree) {  //找到树中最小元素
	if (Tree->left)
		return findMin(Tree->left);
	else
		return Tree;
}
Position findMax(AVLTree Tree) {  //找到树中最大元素
	if (Tree->right)
		return findMin(Tree->left);
	else
		return Tree;
}
int getHeight(AVLTree Tree) { //得到树高
	if (!Tree)return 0;
	return max(getHeight(Tree->left), getHeight(Tree->right)) + 1;
}
int getSubTreeSize(AVLTree Tree) {
	int cnt = 0;
	if (Tree->left)
		cnt = Tree->left->subTreeSize + 1;
	else cnt = 1;
	if (Tree->right)
		cnt += Tree->right->subTreeSize;
	return cnt;
}
AVLTree LL_Rotation(AVLTree Tree) { //LL旋转
	AVLTree tmp = Tree->left;
	Tree->left = tmp->right;
	tmp->right = Tree;
	Tree->height = max(getHeight(Tree->left), getHeight(Tree->right)) + 1;
	tmp->height = max(getHeight(tmp->left), Tree->height) + 1;
	tmp->subTreeSize = Tree->subTreeSize;
	Tree->subTreeSize = getSubTreeSize(Tree);
	return tmp;
}
AVLTree RR_Rotation(AVLTree Tree) {  //RR旋转
	AVLTree tmp = Tree->right;
	Tree->right = tmp->left;
	tmp->left = Tree;
	Tree->height = max(getHeight(Tree->left), getHeight(Tree->right)) + 1;
	tmp->height = max(getHeight(tmp->right), Tree->height) + 1;
	tmp->subTreeSize = Tree->subTreeSize;
	Tree->subTreeSize = getSubTreeSize(Tree);
	return tmp;
}
AVLTree LR_Rotation(AVLTree Tree){  //LR旋转
	Tree->left = RR_Rotation(Tree->left);
	return LL_Rotation(Tree);
}
AVLTree RL_Rotation(AVLTree Tree) { //RL旋转
	Tree->right = LL_Rotation(Tree->right);
	return RR_Rotation(Tree);
}
AVLTree addNewNode(AVLTree Tree, int value){
	if (!Tree) { 
		Tree = (AVLTree)malloc(sizeof(struct AVLNode));
		Tree->value = value;
		Tree->height = 0;
		Tree->left = Tree->right = NULL;
		Tree->subTreeSize = 1;
	} 
	else if (value < Tree->value) { //左子树中
		Tree->left = addNewNode(Tree->left, value);
		Tree->subTreeSize++;
		if (getHeight(Tree->left) - getHeight(Tree->right) == 2)//若平衡被破坏
			if (value < Tree->left->value)  //判断是进行LL旋转还是LR旋转
				Tree = LL_Rotation(Tree);
			else
				Tree = LR_Rotation(Tree);
	}
	else if (value > Tree->value) {//又子树中
		Tree->right = addNewNode(Tree->right, value);
		Tree->subTreeSize++;
		if (getHeight(Tree->left) - getHeight(Tree->right) == -2)//若平衡被破坏
			if (value > Tree->right->value)//判断是进行RR旋转还是RL旋转
				Tree = RR_Rotation(Tree);
			else
				Tree = RL_Rotation(Tree);
	}
	//更新树高
	Tree->height = max(getHeight(Tree->left), getHeight(Tree->right)) + 1;
	return Tree;
}
AVLTree deleteOldNode(AVLTree Tree, int value) { //删除节点
	Position tmp;
	if (!Tree)
		printf("not find\n");
	else {
		if (value < Tree->value) {//往左,删除节点后可能导致右子树失去平衡
			Tree->subTreeSize--;  //则以该节点为根节点子树大小减一
			Tree->left = deleteOldNode(Tree->left, value);
			if (getHeight(Tree->left) - getHeight(Tree->right) == -2) {//若平衡被破坏
				tmp = Tree->right;
				if (getHeight(tmp->left)> getHeight(tmp->right))//判断是进行RR旋转还是RL旋转
					Tree = RL_Rotation(Tree);  //若左子树高度大于右子树则做RL旋转
				else
					Tree = RR_Rotation(Tree);  //否则做RR旋转
			}
		}
		else if (value > Tree->value) {//往右，删除节点后可能导致左子树失去平衡
			Tree->right = deleteOldNode(Tree->right, value);
			Tree->subTreeSize--;
			if (getHeight(Tree->left) - getHeight(Tree->right) == 2) {//若平衡被破坏
				tmp = Tree->left;
				if (getHeight(tmp->left) > getHeight(tmp->right))  //判断是进行LL旋转还是LR旋转
					Tree = LL_Rotation(Tree);//原理同上
				else
					Tree = LR_Rotation(Tree);
			}
		}
		else {//找到要删除节点
			if (Tree->left && Tree->right) {  //若节点左右儿子均存在
				if (getHeight(Tree->left) > getHeight(Tree->right)) { //为尽量使得树平衡，判断去寻找左子树的最大值还是右子树的最小值
					tmp = findMax(Tree->left);	//找到左子树中最大元素
					Tree->value = tmp->value;  //与要删除节点互换位置
					Tree->subTreeSize--;  //更新子树大小
					Tree->left= deleteOldNode(Tree->left, Tree->value);
				}
				else {
					tmp = findMin(Tree->right);  //找到右子树中最小元素
					Tree->value = tmp->value;  //与要删除节点互换位置
					Tree->subTreeSize--;  //更新子树大小
					Tree->right = deleteOldNode(Tree->right, Tree->value); //到右子树中将相应位置删除
				}
			}
			else { //有一个或无儿子
				tmp = Tree;
				if (!Tree->left)       //只有右儿子或无子结点
					Tree = Tree->right;
				else                   // 只有左儿子 
					Tree = Tree->left;
				free(tmp);
			}
		}
	}
	return Tree;
}
AVLTree createTree(int a[], int n) {
	if (!n)return NULL;
	AVLTree Tree = NULL;
	for (int i = 0; i < n; i++)
		Tree = addNewNode(Tree, a[i]);
	return Tree;
}
int findKth(AVLTree Tree, int k) {//寻找第k小的数
	if (k > Tree->subTreeSize)return -inf;
	if (!Tree->left && !Tree->right)return Tree->value;
	if (Tree->left) {
		int now = Tree->left->subTreeSize + 1;
		if (now == k)return Tree->value;
		else if (now > k)
			return findKth(Tree->left, k);
		if (Tree->right)
			return findKth(Tree->right, k - now);
	}
	else {
		if (k == 1)
			return Tree->value;
		else
			return findKth(Tree->right, k - 1);
	}
}

调试代码

void qxbl(AVLTree Tree) { //前序遍历
	if (Tree) {
		printf("%d %d\n", Tree->value, Tree->subTreeSize);
		qxbl(Tree->left);
		qxbl(Tree->right);
	}
}
int main() {
	int n = 8;
	int a[] = { 1,2,9,4,5,11,7,8 };
	AVLTree Tree = createTree(a, n);
	qxbl(Tree);
	char ty;
	int k;
	while (cin >> ty) {  //判断操作类型 a为插入节点 b为删除节点 c为查询第k小的数
		cin >> k;
		if (ty == 'a') {
			Tree = addNewNode(Tree, k);//添加新节点
			printf("Add success!\n");
			qxbl(Tree);
		}
		else if (ty == 'b') {
			Tree = deleteOldNode(Tree, k); //删除节点
			printf("Delete successful!\n");
			qxbl(Tree);
		}
		else
			printf("The %d-th small number:%d\n", k, findKth(Tree, k));
	}
	return 0;
}

运行结果如下

5 8
2 3
1 1
4 1
9 4
7 2
8 1
11 1
a 3
Add success!
5 9
2 4
1 1
4 2
3 1
9 4
7 2
8 1
11 1
c 5
The 5-th small number:5
b 5
Delete successful!
7 8
2 4
1 1
4 2
3 1
9 3
8 1
11 1
c 5
The 5-th small number:7