【BZOJ 3224】普通平衡树

BZOJ 3224

题意

有n个操作，一共有6种。

编号	操作	含义
1	insert(x)	插入一个数x
2	remove(x)	删去一个数x，多个删去一个
3	rank(x)	求出x的排名，多个取最小的
4	kth(x)	求第k小的数
5	pred(x)	小于x中最大的数
6	succ(x)	大于x中最小的数

样例输入

10
1 106465
4 1
1 317721
1 460929
1 644985
1 84185
1 89851
6 81968
1 492737
5 493598

样例输出

106465
84185
492737

SOL

这道题其实把SBT（Size Balanced Tree平衡二叉树）的基本操作都覆盖进去了，下面一个个讲解。
注意：我的方法是对每一个数都建立一个节点，也可以把相同的数放在一个节点上，不过写起来有点复杂，放在最后。

定义

SBT的定义：一个节点，任何一个孩子不小于其他孩子的孩子。
简单来说就是任何一个节点都比它的兄弟的孩子大。
如图：L比C和D大，R比A和B大。

定义结构体tree，其中key为节点值，left和right为左右孩子编号，size为子树大小。

struct SBT
    {
        int key,left,right,size;
    } tree[N];  

旋转

其实旋转才是SBT最基本的操作，这也是为什么SBT能保持平衡性质。
如图所示，基本上看看代码就能看懂了，不要漏了操作。

void left_rotate(int &x)
{  
    int y = tree[x].right;  
    tree[x].right = tree[y].left;  
    tree[y].left = x;  
    tree[y].size = tree[x].size;
    tree[x].size = tree[tree[x].left].size + tree[tree[x].right].size + 1;  
    x = y;  
}  
void right_rotate(int &x)  
{  
    int y = tree[x].left;  
    tree[x].left = tree[y].right;  
    tree[y].right = x;  
    tree[y].size = tree[x].size;  
    tree[x].size = tree[tree[x].left].size + tree[tree[x].right].size + 1;  
    x = y;  
}  

维护

SBT还有一个基本操作就是维护（maintain）。
是啊讲了这么多平衡性质还没说怎么实现呢！

分成以下几种情况：

一、A大于R
1、调用right_rotate(T)；

2、R并没有变化，L和T都发生了变化，所以调用一次maintain(T)和maintain(T)。

二、B大于R

1、调用left_rotate(L)；

2、调用right_rotate(T)；

3、此时R没有发生变化，所以调用maintain(L)，maintain(T)，maintain(B)。

Tip！

注意调用maintain时要从下到上调用！否则上面的旋转会影响下面的节点。

对于另外两种情况同理。

void maintain(int &x,bool flag)  
{  
    if (flag == false)
    {  
        if (tree[tree[tree[x].left].left].size > tree[tree[x].right].size)//左孩子的左子树大于右孩子  
            right_rotate(x);  
        else if (tree[tree[tree[x].left].right].size > tree[tree[x].right].size)//右孩子的右子树大于右孩子  
        {  
            left_rotate(tree[x].left);  
            right_rotate(x);  
        }  
        else return;  
    }  
    else
    {  
        if(tree[tree[tree[x].right].right].size > tree[tree[x].left].size)//右孩子的右子树大于左孩子  
            left_rotate(x);  
        else if(tree[tree[tree[x].right].left].size > tree[tree[x].left].size)//右孩子的左子树大于左孩子  
        {  
            right_rotate(tree[x].right);  
            left_rotate(x);  
        }  
        else return;  
    }  
    maintain(tree[x].left,false);  
    maintain(tree[x].right,true);  
    maintain(x,true);  
    maintain(x,false);  
}

插入

终于进入正题！
插入操作比较简单，只要在插入好后用maintain维护一下即可。

void insert(int &x,int key)  
{  
    if(x == 0)  
    {  
        x = ++tot;  
        tree[x].left = tree[x].right = 0;  
        tree[x].size = 1;  
        tree[x].key = key;  
    }  
    else  
    {  
        tree[x].size ++;  
        if(key < tree[x].key) insert(tree[x].left,key);  
        else insert(tree[x].right,key);
        maintain(x,key>=tree[x].key); 
    }  
}

删除

删除就有些麻烦了，我这里的方法是，如果没有子树那最好；如果只有左子树或者只有右子树，就把子树挂上去；如果有两棵子树那么找到后继节点代替当前位置。其实讲道理要调用一下maintain但是其实无所谓这么一点。。。

void remove(int &x,int key)   
{  
    tree[x].size --;  
    if (key > tree[x].key)  
        remove(tree[x].right,key);  
    else if(key < tree[x].key)  
        remove(tree[x].left,key);  
    else  
    {  
        if(tree[x].left != 0 && tree[x].right == 0) x = tree[x].left;//只有左子树
        else if(tree[x].left ==0 && tree[x].right != 0) x = tree[x].right;//只有右子树
        else if(tree[x].left ==0 && tree[x].right == 0) x = 0;//无左子树和右子树
        else //找到右子树中最小元素
        {  
            int temp = tree[x].right;  
            while(tree[temp].left) temp = tree[temp].left;  
            tree[x].key = tree[temp].key;
            remove(tree[x].right,tree[temp].key);  
        }  
    }  
}

求排名

不断的在二叉树上搜索。这里就是重复开辟节点的唯一的麻烦：由于节点可能有重复，左子树可能还有相同的节点，所以相等了也不要马上停止，而是记录下当期答案继续向左搜索。

void GetRank(int &x,int key,int sum)//求key排第几  
{  
    if (x == 0) return;
    if (key == tree[x].key)
        {
            res = min(res,sum+tree[tree[x].left].size+1);
            GetRank(tree[x].left,key,sum);
        }
    if(key < tree[x].key) GetRank(tree[x].left,key,sum);  
    if(key > tree[x].key) GetRank(tree[x].right,key,sum+tree[tree[x].left].size+1); 
}  

求第k小值

重复开节点求第k小值相当方便，左子树大小严格等于k-1即可。

int GetKth(int &x,int k)//求第k小数  
{  
    int r = tree[tree[x].left].size + 1;  
    if(r == k) return tree[x].key;  
    else if(r < k) return GetKth(tree[x].right,k - r);  
    else return GetKth(tree[x].left,k);  
} 

前驱和后继

int pred(int &x,int y,int key)
{  
    if(x == 0) return y;  
    if(key > tree[x].key)  
        return pred(tree[x].right,x,key);  
    else return pred(tree[x].left,y,key);  
}  

int succ(int &x,int y,int key)
{  
    if(x == 0) return y;  
    if(key < tree[x].key)  
        return succ(tree[x].left,x,key);  
    else return succ(tree[x].right,y,key);  
} 

更好的写法

丢个链接 http://www.cnblogs.com/BeyondW/p/5719597.html。
这位大神用的是压缩写法，可以发现在rotate和maintain操作中会有大量的对称重复情况，用0和1表示左右孩子代码量少很多，但是看起来就不那么容易。

时间复杂度分析

接下来就是最重要的效率分析啦！
其实我开始会觉得maintain操作复杂度挺高的，怎么说也得要接近log吧，但是事实上均摊复杂度为O(1),Orz我怎么看不出。。。
由于树方面的证明过于高深，所以还是 ：
http://www.cnblogs.com/zgmf_x20a/archive/2008/11/14/1333205.html。
这里有对树高的分析和maintain操作代价的分析，不想看的：
树高稳定在log(n)，maintain均摊复杂度O(1)，所以不管带不带maintain操作，所有操作时间复杂度均为log(n)！

完整代码

重复数字开节点代码

#include<cmath>
#include<cstdio>
#include<vector>
#include<cstring>
#include<iomanip>
#include<stdlib.h>
#include<iostream>
#include<algorithm>
#define ll long long
#define inf 1000000000
#define mod 1000000007
#define N 5000000
using namespace std;
int n,i,ch;
struct t{int key,left,right,size;} tree[N];  
int t1,t2,root,tot,sum,res;  
void left_rotate(int &x)
{  
    int y = tree[x].right;  
    tree[x].right = tree[y].left;  
    tree[y].left = x;  
    tree[y].size = tree[x].size;
    tree[x].size = tree[tree[x].left].size + tree[tree[x].right].size + 1;  
    x = y;  
}  
void right_rotate(int &x)  
{  
    int y = tree[x].left;  
    tree[x].left = tree[y].right;  
    tree[y].right = x;  
    tree[y].size = tree[x].size;  
    tree[x].size = tree[tree[x].left].size + tree[tree[x].right].size + 1;  
    x = y;  
}  
void maintain(int &x,bool flag)  
{  
    if (flag == false)
    {  
        if (tree[tree[tree[x].left].left].size > tree[tree[x].right].size)//左孩子的左子树大于右孩子  
            right_rotate(x);  
        else if (tree[tree[tree[x].left].right].size > tree[tree[x].right].size)//右孩子的右子树大于右孩子  
        {  
            left_rotate(tree[x].left);  
            right_rotate(x);  
        }  
        else return;  
    }  
    else
    {  
        if(tree[tree[tree[x].right].right].size > tree[tree[x].left].size)//右孩子的右子树大于左孩子  
            left_rotate(x);  
        else if(tree[tree[tree[x].right].left].size > tree[tree[x].left].size)//右孩子的左子树大于左孩子  
        {  
            right_rotate(tree[x].right);  
            left_rotate(x);  
        }  
        else return;  
    }  
    maintain(tree[x].left,false);  
    maintain(tree[x].right,true);  
    maintain(x,true);  
    maintain(x,false);  
}
void insert(int &x,int key)  
{  
    if(x == 0)  
    {  
        x = ++tot;  
        tree[x].left = tree[x].right = 0;  
        tree[x].size = 1;  
        tree[x].key = key;  
    }  
    else  
    {  
        tree[x].size ++;  
        if(key < tree[x].key) insert(tree[x].left,key);  
        else insert(tree[x].right,key);
        maintain(x,key>=tree[x].key); 
    }  
}
void remove(int &x,int key)   
{  
    tree[x].size --;  
    if (key > tree[x].key)  
        remove(tree[x].right,key);  
    else if(key < tree[x].key)  
        remove(tree[x].left,key);  
    else  
    {  
        if(tree[x].left != 0 && tree[x].right == 0) x = tree[x].left;//只有左子树
        else if(tree[x].left ==0 && tree[x].right != 0) x = tree[x].right;//只有右子树
        else if(tree[x].left ==0 && tree[x].right == 0) x = 0;//无左子树和右子树
        else //找到右子树中最小元素
        {  
            int temp = tree[x].right;  
            while(tree[temp].left) temp = tree[temp].left;  
            tree[x].key = tree[temp].key;
            remove(tree[x].right,tree[temp].key);  
        }  
    }  
}
int GetKth(int &x,int k)//求第k小数  
{  
    int r = tree[tree[x].left].size + 1;  
    if(r == k) return tree[x].key;  
    else if(r < k) return GetKth(tree[x].right,k - r);  
    else return GetKth(tree[x].left,k);  
} 

void GetRank(int &x,int key,int sum)//求key排第几  
{  
    if (x == 0) return;
    if (key == tree[x].key)
        {
            res = min(res,sum+tree[tree[x].left].size+1);
            GetRank(tree[x].left,key,sum);
        }
    if(key < tree[x].key) GetRank(tree[x].left,key,sum);  
    if(key > tree[x].key) GetRank(tree[x].right,key,sum+tree[tree[x].left].size+1); 
}  

int pred(int &x,int y,int key)
{  
    if(x == 0) return y;  
    if(key > tree[x].key)  
        return pred(tree[x].right,x,key);  
    else return pred(tree[x].left,y,key);  
}  

int succ(int &x,int y,int key)
{  
    if(x == 0) return y;  
    if(key < tree[x].key)  
        return succ(tree[x].left,x,key);  
    else return succ(tree[x].right,y,key);  
}  

int main()  
{
    root = tot = 0;   
    scanf("%d",&n);
    int x,tmp;  
    for (i = 1;i <= n; i++)
    {  
        scanf("%d%d",&ch,&x);
        if (ch==1) insert(root,x);
        if (ch==2) remove(root,x);
        if (ch==3) {res=inf;GetRank(root,x,0);printf("%d\n",res);}
        if (ch==4) printf("%d\n",GetKth(root,x));
        if (ch==5) printf("%d\n",tree[pred(root,0,x)].key);
        if (ch==6) printf("%d\n",tree[succ(root,0,x)].key);
    }  
    return 0;  
}  

重复数字不开新节点代码

总的来说不开新节点肯定是会快一点，但是代码有些地方乱了很多，会用注释标出，大家可以对照着看。

#include<cmath>
#include<cstdio>
#include<vector>
#include<cstring>
#include<iomanip>
#include<stdlib.h>
#include<iostream>
#include<algorithm>
#define ll long long
#define inf 1000000000
#define mod 1000000007
#define N 5000000 
int n,i,ch;
struct SBT
{  
    int key,left,right,size;  
} tree[N];  

int t1,t2,root,tot,sum,res,flag;  

void left_rotate(int &x)  
{  
    int y = tree[x].right;  
    int sx = tree[x].size - tree[tree[x].left].size - tree[tree[x].right].size;//旋转的时候每个节点不再只有1个数了，要重新记录
    tree[x].right = tree[y].left;  
    tree[y].left = x;  
    tree[y].size = tree[x].size;
    tree[x].size = tree[tree[x].left].size + tree[tree[x].right].size + sx;  
    x = y;  
}  

void right_rotate(int &x)  
{  
    int y = tree[x].left;  
    int sx = tree[x].size - tree[tree[x].left].size - tree[tree[x].right].size;//同上
    tree[x].left = tree[y].right;  
    tree[y].right = x;  
    tree[y].size = tree[x].size;  
    tree[x].size = tree[tree[x].left].size + tree[tree[x].right].size + sx;  
    x = y;  
}  

void maintain(int &x,bool flag)  
{  
    if(flag == false)
    {  
        if(tree[tree[tree[x].left].left].size > tree[tree[x].right].size)
            right_rotate(x);  
        else if(tree[tree[tree[x].left].right].size > tree[tree[x].right].size)
        {  
            left_rotate(tree[x].left);  
            right_rotate(x);  
        }  
        else return;  
    }  
    else
    {  
        if(tree[tree[tree[x].right].right].size > tree[tree[x].left].size)
            left_rotate(x);  
        else if(tree[tree[tree[x].right].left].size > tree[tree[x].left].size)
        {  
            right_rotate(tree[x].right);  
            left_rotate(x);  
        }  
        else return;  
    }  
    maintain(tree[x].left,false);  
    maintain(tree[x].right,true);  
    maintain(x,true);  
    maintain(x,false);  
}  

void insert(int &x,int key)  
{  
    if(x == 0)  
    {  
        x = ++tot;  
        tree[x].left = tree[x].right = 0;  
        tree[x].size = 1;  
        tree[x].key = key;  
    }  
    else  
    {  
        tree[x].size ++;
        if (key < tree[x].key) insert(tree[x].left,key);  
        if (key > tree[x].key) insert(tree[x].right,key);//相等的时候就不用再递归了
        maintain(x, key >= tree[x].key);
    }  
}  


void remove(int &x,int key)   
{  
    tree[x].size --;  
    if(key > tree[x].key)  
        remove(tree[x].right,key);  
    else if(key < tree[x].key)  
        remove(tree[x].left,key);  
    else if (tree[x].size - tree[tree[x].left].size - tree[tree[x].right].size > 0) return;//如果当前节点还存在就不用删除了
    else
    {  
        if(tree[x].left != 0 && tree[x].right == 0) x = tree[x].left;
        else if(tree[x].right !=0 && tree[x].left == 0) x = tree[x].right;
        else if(!tree[x].left && !tree[x].right) x = 0;
        else 
            {
                right_rotate(x);
                int h = tree[x].right;
                tree[h].size = tree[tree[h].left].size + tree[tree[h].right].size + 1;
                remove(tree[x].right,tree[tree[x].right].key);
                maintain(x,true);
            } 
    }  
}  


int GetKth(int &x,int k)
{  
    int s1 = tree[tree[x].left].size;
    int s2 = tree[x].size - tree[tree[x].left].size - tree[tree[x].right].size;
    if (s1+1<=k&&k<=s1+s2) return tree[x].key;//这里有可能当前节点“包含”了第k个，所以不能简单地等于
    if(k > s1+s2) return GetKth(tree[x].right,k - s1 - s2);  
    else return GetKth(tree[x].left,k);  
} 

int GetRank(int &x,int key)//这里就直接左子树大小+1即可，唯一稍微简单的地方
{  
    int s1 = tree[tree[x].left].size;
    int s2 = tree[x].size - tree[tree[x].left].size - tree[tree[x].right].size;
    if (key == tree[x].key) 
        if (tree[x].left == 0) return 1; else return s1+1;
    if(key < tree[x].key) return GetRank(tree[x].left,key);  
    if(key > tree[x].key) return GetRank(tree[x].right,key)+s1+s2; 
}  

int pred(int &x,int y,int key)
{  
    if(x == 0) return y;  
    if(tree[x].key < key)  
        return pred(tree[x].right,x,key);  
    else return pred(tree[x].left,y,key);  
}  

int succ(int &x,int y,int key)
{  
    if(x == 0) return y;  
    if(tree[x].key > key)  
        return succ(tree[x].left,x,key);  
    else return succ(tree[x].right,y,key);  
}  

int main()  
{
    root = tot = 0;   
    scanf("%d",&n);
    int x,tmp;  
    for (i = 1;i <= n; i++)
    {  
        scanf("%d%d",&ch,&x);
        if (ch==1) insert(root,x);
        if (ch==2) remove(root,x);
        if (ch==3) printf("%d\n",GetRank(root,x));
        if (ch==4) printf("%d\n",GetKth(root,x));
        if (ch==5) printf("%d\n",tree[pred(root,0,x)].key);
        if (ch==6) printf("%d\n",tree[succ(root,0,x)].key);
    }  
    return 0;  
}