纯萌新,学了个avl树想来试试这道题,可是题解也没有用这个的,忙活半天也找不到原因。
#include <iostream>
using namespace std;
int k=0;
struct TreeNode {
int val{}; // 节点值
int cnt=0; //计数
int lsum=0,rsum=0;
int height = 0; // 节点高度
TreeNode *left{}; // 左子节点
TreeNode *right{}; // 右子节点
TreeNode() = default;
explicit TreeNode(int x) : val(x),cnt(1){}
};
inline int get_height(TreeNode* node)
{
return node==nullptr?-1:node->height;
}
void updateHeight(TreeNode* node)
{
node->height=max(get_height(node->left),get_height(node->right))+1;
}
inline int get_balanceFactor(TreeNode* node)
{
if(node==nullptr)
return 0;
return get_height(node->left)-get_height(node->right);
}
/* 左旋操作 */
TreeNode *leftRotate(TreeNode *node) {
TreeNode *child = node->right;
TreeNode *grandChild = child->left;
node->rsum = grandChild ? grandChild->cnt + grandChild->lsum + grandChild->rsum : 0;
child->lsum+=node->lsum+node->cnt;
// 以 child 为原点,将 node 向左旋转
child->left = node;
node->right = grandChild;
// 更新节点高度
updateHeight(node);
updateHeight(child);
// 返回旋转后子树的根节点
return child;
}
/* 右旋操作 */
TreeNode *rightRotate(TreeNode *node) {
TreeNode *child = node->left;
TreeNode *grandChild = child->right;
node->lsum = grandChild ? grandChild->cnt + grandChild->lsum + grandChild->rsum : 0;
child->rsum+=node->rsum+node->cnt;
// 以 child 为原点,将 node 向右旋转
child->right = node;
node->left = grandChild;
// 更新节点高度
updateHeight(node);
updateHeight(child);
// 返回旋转后子树的根节点
return child;
}
TreeNode* rotate(TreeNode* node)
{
int _balanceFactor=get_balanceFactor(node);
//左偏树
if (_balanceFactor > 1) {
if (get_balanceFactor(node->left) >= 0) {
// 右旋
return rightRotate(node);
} else {
// 先左旋后右旋
node->left = leftRotate(node->left);
return rightRotate(node);
}
}
// 右偏树
if (_balanceFactor < -1) {
if (get_balanceFactor(node->right) <= 0) {
// 左旋
return leftRotate(node);
} else {
// 先右旋后左旋
node->right = rightRotate(node->right);
return leftRotate(node);
}
}
// 平衡树,无须旋转,直接返回
return node;
}
/* 递归插入节点(辅助方法) */
TreeNode* insertHelper(TreeNode *node, int val)
{
if (node == nullptr)
return new TreeNode(val);
/* 1. 查找插入位置并插入节点 */
if (val < node->val)
{
node->left = insertHelper(node->left, val);
node->lsum++;
}
else if (val > node->val)
{
node->right = insertHelper(node->right, val);
node->rsum++;
}
else
node->cnt++;
updateHeight(node); // 更新节点高度
/* 2. 执行旋转操作,使该子树重新恢复平衡 */
node = rotate(node);
node->lsum = node->left ? node->left->lsum + node->left->cnt+node->left->rsum : 0;
node->rsum = node->right ? node->right->rsum + node->right->cnt+node->right->lsum: 0;
// 返回子树的根节点
return node;
}
/* 插入节点 */
void insert(int val,TreeNode*& root) {
root =insertHelper(root, val);
}
/* 递归删除节点(辅助方法) */
TreeNode *removeHelper(TreeNode *node, int val,int cnt_to_remove = 1) {
if (node == nullptr)
return nullptr;
/* 1. 查找节点并删除 */
if (val < node->val)
{
// 递归去左子树删除
int before = node->lsum; // 保存原来的 lsum
node->left = removeHelper(node->left, val,cnt_to_remove);
// 如果左子树的 lsum 发生了变化,则说明实际上有节点被删除了
if (before > (node->left ? node->left->lsum + node->left->cnt+node->left->rsum : 0)) {
node->lsum-=cnt_to_remove; // 实际上删除了节点,才减少 lsum
}
}
else if (val > node->val)
{
int before = node->rsum; // 保存原来的 rsum
node->right = removeHelper(node->right, val,cnt_to_remove);
// 如果右子树的 rsum 发生了变化,则说明实际上有节点被删除了
if (before > (node->right ? node->right->rsum + node->right->cnt+node->right->lsum : 0)) {
node->rsum-=cnt_to_remove; // 实际上删除了节点,才减少 rsum
}
}
else {
if(node->cnt>cnt_to_remove)
{
node->cnt-=cnt_to_remove;
return node;
}
if (node->left == nullptr || node->right == nullptr) {
TreeNode *child = node->left != nullptr ? node->left : node->right;
// 子节点数量 = 0 ,直接删除 node 并返回
if (child == nullptr) {
delete node;
return nullptr;
}
// 子节点数量 = 1 ,直接删除 node
else {
delete node;
node = child;
}
} else {
// 子节点数量 = 2 ,则将中序遍历的下个节点删除,并用该节点替换当前节点
TreeNode *temp = node->right;
while (temp->left != nullptr) {
temp = temp->left;
}
int tempVal = temp->val;
int tempCnt=temp->cnt;
node->val = tempVal;
node->cnt=tempCnt;
node->right = removeHelper(node->right, temp->val,tempCnt);
if(node->right)
node->rsum=node->right->rsum+node->right->lsum+node->right->cnt;
else
{
node->rsum=0;
}
}
}
updateHeight(node); // 更新节点高度
/* 2. 执行旋转操作,使该子树重新恢复平衡 */
node = rotate(node);
node->lsum = node->left ? node->left->lsum + node->left->cnt+node->left->rsum : 0;
node->rsum = node->right ? node->right->rsum + node->right->cnt+node->right->lsum: 0;
// 返回子树的根节点
return node;
}
/* 删除节点 */
void remove(int val,TreeNode*& root) {
root = removeHelper(root, val);
}
int search(TreeNode* root, int val, int& k) {
if (!root) return k+1;
if (root->val < val) {
// 当前节点的值小于目标值,增加k值(左子树节点数量 + 当前节点)
k += root->lsum + root->cnt;
return search(root->right, val, k); // 继续在右子树中搜索
} else if (root->val > val) {
return search(root->left, val, k); // 继续在左子树中搜索
} else {
// 找到目标值,返回比它小的所有数的数量 + 1
return root->lsum + k + 1;
}
}
void rank1(int rank_,TreeNode* root)
{
if(!root)return;
int tmp=root->lsum;
if((tmp+1)==rank_)
{
cout<<root->val<<endl;
return;
}
if((tmp+1)>rank_)
{
rank1(rank_,root->left);
}
else
{
rank1(rank_-tmp-root->cnt,root->right);
}
}
void leftsearch(int val, TreeNode*root)
{
k=0;
rank1(search(root,val,k)-1,root);
}
void rightsearch(int val, TreeNode*root)
{
k=0;
rank1(search(root,val+1,k),root);
}
int main()
{
TreeNode* root=nullptr;
int n;
cin>>n;
int opt;
int temp;
while(n--)
{
cin>>opt;
if(opt==1)
{
cin>>temp;
insert(temp,root);
}
else if(opt==2)
{
cin>>temp;
remove(temp,root);
}
else if(opt==3)
{
k=0;
cin>>temp;
cout<<search(root,temp,k)<<endl;
}
else if(opt==4)
{
cin>>temp;
rank1(temp, root);
}
else if(opt==5)
{
cin>>temp;
leftsearch(temp,root);
}
else if(opt==6)
{
cin>>temp;
rightsearch(temp, root);
}
}
}