用LUA（和C++）刷PAT (Advanced Level) ——1066 Root of AVL Tree

最新推荐文章于 2024-11-14 16:58:13 发布

起变量名什么的简直神烦

最新推荐文章于 2024-11-14 16:58:13 发布

阅读量90

点赞数

文章标签： lua c++

本文链接：https://blog.csdn.net/weixin_44316978/article/details/120043412

版权

链表这种东西，真的很适合一个bug查一天…

#include <iostream>
#include <algorithm>
#include <cstring>

using namespace std;

struct node {
	int key;
	node *parent, *left, *right;
	int balance = 0, height = 0;
};

node* root = NULL;

void insert(node* root_node, node* node_add) {
	if (root_node->key < node_add->key) {
		if (root_node->right == NULL) {
			root_node->right = node_add;
			node_add->parent = root_node;
		}
		else
			insert(root_node->right, node_add);
	}
	else {
		if (root_node->left == NULL) {
			root_node->left = node_add;
			node_add->parent = root_node;
		}
		else
			insert(root_node->left, node_add);
	}
}

int getHeight(node* n) { return (n == NULL) ? 0 : n->height; }

void updateData(node* n) {
	n->balance = getHeight(n->left) - getHeight(n->right);
	n->height = max(getHeight(n->left), getHeight(n->right)) + 1;
}

void leftRotate(node* n)
{
	node* n_right = n->right;
	node* n_parent = n->parent;

	n->right = n_right->left;
	if (n->right != NULL)
		n->right->parent = n;

	n_right->left = n;
	n->parent = n_right;

	if (n_parent == NULL) {
		root = n_right;
		root->parent = NULL;
	}
	else {
		n_right->parent = n_parent;
		if (n_parent->left == n)
			n_parent->left = n_right;
		else
			n_parent->right = n_right;
	}
	updateData(n);
	updateData(n_right);
}

void rightRotate(node* n)
{
	node* n_left = n->left;
	node* n_parent = n->parent;

	n->left = n_left->right;
	if (n->right != NULL)
		n->right->parent = n;

	n_left->right = n;
	n->parent = n_left;

	if (n_parent == NULL) {
		root = n_left;
		root->parent = NULL;
	}
	else {
		n_left->parent = n_parent;
		if (n_parent->left == n)
			n_parent->left = n_left;
		else
			n_parent->right = n_left;
	}
	updateData(n);
	updateData(n_left);
}

void processUnbalance(node* n) {
	if (n->balance > 1 && n->left->balance > 0)
		rightRotate(n);
	else if (n->balance > 1 && n->left->balance < 0) {
		leftRotate(n->left);
		rightRotate(n);
	}
	else if (n->balance < -1 && n->right->balance < 0)
		leftRotate(n);
	else if (n->balance < -1 && n->right->balance > 0) {
		rightRotate(n->right);
		leftRotate(n);
	}
}

void updateInfo(node* n) {
	updateData(n);
	if (n->balance > 1 || n->balance < -1)
		processUnbalance(n);
	if (n->parent != NULL)
		updateInfo(n->parent);
}

int main()
{
	int N;
	cin >> N;
	for (int i = 0; i < N; i++) {
		node* n = (node*)malloc(sizeof(node));
		cin >> n->key;
		n->parent = NULL;
		n->right = NULL;
		n->left = NULL;
		if (root == NULL)
			root = n;
		else
		{
			insert(root, n);
			updateInfo(n);
		}
	}
	cout << root->key;
}