文章目录
- 写在前面
- 红黑树是什么
- 红黑树的平衡性
- 红黑树整体框架
- 旋转操作
- 插入操作
- 双红修正
- Case 1, 2
- Case 3
- Case 4
- Case 5
- Case 6
- 删除操作
- 二叉查找树
- 红黑树
- Case 0
- Case 1
- Case 2
- Case 3
- 双黑修正
- Case 1
- Case 2
- Case 3
- Case 4
- Case 5
- 其他查询操作
- 查询排名
- 查询第k大
- 寻找前驱
- 寻找后继
- 最终代码
写在前面
推荐一个很实用的工具:红黑树可视化
本文参考OI wiki中的红黑树代码,读者也可以参考该篇解析(写得还是很不错的),不过OI Wiki里删除后平衡维护的Case 4和Case 5在代码细节上稍微有些问题(把 c c c, d d d 均为红色算进Case 4了,这样不会出bug,只是相当于绕了个弯)。
大部分红黑树代码都采用 rotateLeft 和 rotateRight 两个函数来进行旋转,而且在找close/distant nephew的时候也是分类讨论,这样比较麻烦。
我们其实可以使用 node *ch[2] 来表示左右孩子,ch[0] 表示左孩子,ch[1] 表示右孩子。在后续使用中,我们#define left ch[0], #define right ch[1],从而兼容用 left, right 找左右孩子的方法。这种存储方式的好处是,我们可以通过异或1来轻松切换左右分支,而不是采用三目运算符来找兄弟节点。
再考虑 rotate,其实 rotate 相当于把某个子节点往上转到其父节点的位置,因此我们可以用同一个 rotate 函数来表示左旋或右旋,使用时旋转左孩子即为 rotateLeft,旋转右孩子即为 rotateRight。
红黑树是什么
如果能把一个又一个节点积累起来,也许就能变成一棵红黑树。 ——高松灯
红黑树是一棵满足特殊性质的二叉搜索树。它的特殊性质有:
- 节点均为红色/黑色(顾名思义)
- NIL 节点(空叶子节点,有时也叫外部节点)为黑色
- 红色节点的子节点均为黑色
- 从根节点(不含根节点本身)到 NIL 节点的每条路径上的黑色节点数量相同
(注:有的教材/解析要求根节点一定是黑色,不过这个没有太大影响,根节点是红色也不会影响树的平衡性)
我们把第 4 条性质中,从某节点出发(不含节点本身),到任意 NIL 节点路径上的黑节点数,称为节点的“黑高”(black height)。
这里黑高的定义是比较反直觉的(究竟是谁定义的?),因为既要去掉节点本身,又要加上一个额外的 NIL 节点。
可以用以下递推式来加深对黑高的理解:
b
h
[
N
I
L
]
=
0
bh[NIL] = 0
bh[NIL]=0
∀
s
∈
s
o
n
(
x
)
,
\forall s \in son(x),
∀s∈son(x),
b
h
[
x
]
=
b
h
[
s
]
+
[
c
o
l
o
r
(
s
)
=
=
B
L
A
C
K
]
bh[x] = bh[s] + [color(s) == BLACK]
bh[x]=bh[s]+[color(s)==BLACK],中括号表示该表达式的bool值,若表达式成立则为1,表达式不成立则为0.
我们把粉色头发的ano酱作为红节点,黑色头发的Rikki作为黑节点,那么下面这棵树是一棵红黑树:
其中蓝色数字标注的是节点的黑高。
假如给上面这棵树的节点填入权值,就是一棵标准的红黑树:
红黑树的平衡性
考虑这样一个问题:一棵“黑高”为
h
h
h的红黑树,至少会有多少节点?(“黑高”是指根到任意NIL节点路径上的黑节点数)
其实可以用归纳法证明结果是
2
h
−
1
2^h-1
2h−1,这里提供另一种思路:
对于黑高确定的红黑树,我们要想节点数最少,直接全部放黑节点就可以了。因为往里面放红节点不会对黑高产生任何影响,而且由于红节点的子节点必须是黑节点,放红节点还会让总节点数变得更多。
而全为黑节点的红黑树必然是一棵满二叉树,因为根节点到任意NIL节点路径上的黑节点数相同,如果某个部分缺了一块,那这部分的黑高就会减少。
那么,黑高为 h h h的红黑树,节点最少的情况是一棵高为 h h h,全是黑节点的满二叉树(注意高是 h h h而不是 h + 1 h+1 h+1,因为计算黑高时多统计了一次NIL节点,又少统计一次自身,二者抵消了)。这样的树有 2 h − 1 2^h-1 2h−1个节点。
因此,一棵“黑高”为 h h h的红黑树,至少有 2 h − 1 2^h-1 2h−1个节点。( N = s i z e o f ( x ) ≥ 2 b h ( x ) − 1 N = \rm sizeof(x)\ge 2^{bh(x)}-1 N=sizeof(x)≥2bh(x)−1)
我们还知道,红节点的两个孩子一定为黑节点,所以说,根节点到NIL节点的路径上,一个红节点就必定有一个黑节点与之对应,所以红节点数一定不超过黑节点数。
因此,
b
h
(
T
r
e
e
)
≥
h
(
T
r
e
e
)
/
2
\rm bh(Tree)\ge h(Tree)/2
bh(Tree)≥h(Tree)/2。
根据上面两个不等式,可以推出:
一棵有 N N N个节点的红黑树,树高不超过 2 log ( N + 1 ) 2\log (N+1) 2log(N+1)。
也就是说,红黑树天生就是平衡的,树高在
O
(
log
N
)
O(\log N)
O(logN)级别。
假如我们能在插入和删除时维护好红黑树的几条性质,我们就能得到一棵高恒为
O
(
log
N
)
O(\log N)
O(logN)的二叉搜索树。
红黑树整体框架
红黑树类定义为:
template<typename T>
class RedBlackTree {
private:
struct node;
node* root; //根节点
//...
public:
int size();
void insert(T);
bool remove(T);
//...
};
节点定义为:
#define RED 1
#define BLACK 0
#define left ch[0]
#define right ch[1]
template<typename T>
struct RedBlackTree<T>::node {
T val; //权值
bool color; //1 is red, 0 is black
node *father, *ch[2];
int siz; //子树大小
int direction() {
if(father == NULL) return 0;
return father->right == this;
}
node* sibling() {
if(father == NULL) return NULL;
return father->ch[direction() ^ 1];
}
node* uncle() {
if(father == NULL) return NULL;
return father->sibling();
}
void pushup() {
siz = (left?left->siz:0) + (right?right->siz:0) + 1;
}
//......
}
其中val代表节点权值,siz代表子树大小,ch[0] 和 ch[1] 分别代表左右孩子。
direction() 表示当前节点所在分支(0为左孩子,1为右孩子),sibling(), uncle() 是在插入/删除中需要经常用到的亲戚节点。为了方便,我们统一提前写好。
旋转操作
可以参考splay树的旋转操作。这里我们不需要区分左旋和右旋,rotate(x) 表示把节点 x x x 旋转到它父亲的位置。
template<typename T>
void RedBlackTree<T>::connect(node *x, node *fa, int k) {
if(x != NULL) x->father = fa;
if(fa != NULL) {
fa->ch[k] = x;
} else {
root = x;
}
}
template<typename T>
void RedBlackTree<T>::rotate(node *x) {
//rotate x to its parent's position
node* y = x->father;
node* z = y->father;
int yson = x->direction(); //the branch of x, 0 is left, 1 is right
if(z == NULL) {
root = x;
x->father = NULL;
} else {
int zson = y->direction();
connect(x, z, zson);
}
connect(x->ch[yson^1], y, yson);
connect(y, x, yson^1);
y->pushup();
x->pushup();
}
插入操作
从今天开始,我们就是一起演奏音乐的命运共同体! ——丰川祥子
红黑树的插入与普通的 BST 的插入操作类似。
我们将新节点作为红节点插入到树中对应位置,再根据相关节点状态进行调整,使整棵树满足红黑树的性质。
具体地说,红黑树要求红节点的子节点均为黑节点,而插入一个红节点可能会使父节点和子节点均为红色,所以在插入后,我们需要进行双红修正。
插入操作的代码实现如下:
template<typename T>
void RedBlackTree<T>::insert(T v) {
node *x = root, *fa = NULL;
while(x != NULL) {
x->siz++;
fa = x;
if(v < x->val) {
x = x->left;
} else {
x = x->right;
}
}
x = new node(v, RED, fa); //create a new node
if(fa == NULL) {
root = x;
} else if(v < fa->val) {
fa->left = x;
} else {
fa->right = x;
}
SolveDoubleRed(x);
}
双红修正
不过,为了下次不失败而努力不就好了?就算失败一次,也要有重来的信心。 ——千早爱音
上面的插入过程中,可能会出现父节点和子节点都是红色的连续双红情况,这违反了红黑树的性质。
在 SolveDoubleRed(x) 函数中,
x
x
x 为红节点,我们检查
x
x
x 的父节点是否为红节点,如果是,则进行修正。(也就是说,这里
x
x
x 表示连续双红节点的子节点)
修正时,我们要保证红黑树的性质成立,即不出现连续的双红节点,以及保证黑高相同。
在下面的所有注释中,我们用<X>
来表示红节点,[X]
表示黑节点,{X}
表示任意颜色的节点。
Case 1, 2
x x x 为根(父节点为空),或 x x x 的父节点为黑,此时无需修正。下面都是需要修正的情况。
Case 3
x x x 的父节点 p p p 为根节点,此时把 p p p 染黑即可。
if(p == root) {
// Case 3: Parent is root and is RED
// Paint parent to BLACK.
// <P> [P]
// | ====> |
// <X> <X>
p->color = BLACK;
return;
}
Case 4
x
x
x (图中的 N)的父节点
p
p
p,叔节点
u
u
u 均为红色。由于该树原本是一棵合法的红黑树,所以
x
x
x 的祖父节点
g
g
g 一定是黑色。
此时我们将
p
p
p 和
u
u
u 染黑,将
g
g
g 染红,这样在
g
g
g 以下就不会有连续的红节点。
由于插入前
g
g
g 到 NIL 节点经历的黑节点数都相同,所以把
p
p
p,
u
u
u 都染黑后黑节点数仍然相同。且因为又把
g
g
g 染为了红色,所以不会对
g
g
g 往上的节点的黑高产生影响。
不过这时节点
g
g
g 与
g
g
g 的父亲又有可能是连续的红节点,因此我们递归对
g
g
g 进行双红修正。
if(x->hasUncle() && x->uncle()->color == RED) {
// Case 4: Both parent and uncle are RED
// Paint parent and uncle to BLACK;
// Paint grandparent to RED;
// Maintain grandparent recursively.
// [G] <G>
// / \ / \
// <P> <U> ====> [P] [U]
// / /
// <X> <X>
p->color = BLACK; //parent -> black
x->uncle()->color = BLACK; //uncle -> black
p->father->color = RED; //grandparent -> red
SolveDoubleRed(p->father);
return;
}
Case 5
叔节点为黑色,且当前节点
x
x
x 与父节点
p
p
p 的方向相反(即一个为左子节点,一个为右子节点)。
这种情况无法直接维护,我们把节点
x
x
x 旋转上来,然后就转化为了Case 6。
把
x
x
x 转上来后,
x
x
x 与
p
p
p 的方向就一致了,这时
p
p
p 成为了
x
x
x 的父节点,于是我们需要处理的节点变成了
p
p
p。
// Case 5 & 6: parent is RED, uncle is BLACK(or NULL)
if(x->direction() != p->direction()) {
// Case 5: Current node is the opposite direction as parent
// Step 1. Rotate x to parent's position.
// Step 2. Goto Case 6.
// [G] [G]
// / \ rotate(X) / \
// <P> [U] ========> <X> [U]
// \ /
// <X> <P>
rotate(x);
x = p; //Now P is the child of double red.
p = x->father; //reset p to x's father
}
Case 6
叔节点为黑色,且当前节点
x
x
x 与父节点
p
p
p 的方向相同(即同为左子节点或右子节点)。
由于父节点
p
p
p 是红色,所以祖父节点
g
g
g 一定是黑色。这种情况下,我们先向上转一次
p
p
p,再把
p
p
p 染黑,
g
g
g 染红。
这里可以证明,这样操作后,树的黑高是不变的。我们假设最左边的图中
u
u
u 的黑高是1(即两个子节点均为 NIL),那么
g
g
g 的黑高是2,则由于黑高相同的性质,
p
p
p 和
x
x
x(图中的N)黑高也为2.
在旋转+重新染色后,
x
x
x 和
u
u
u 的子树结构和颜色没有变化,因此
x
x
x 的黑高仍为2,
u
u
u 的黑高仍为1. 那么
g
g
g 的黑高为2,
p
p
p 的黑高为2,与开始时
g
g
g 的黑高一致。
// Case 6: Current node is the same direction as parent
// Step 1. Rotate parent to grandparent's position
// Step 2. Paint parent (before rotate) to BLACK;
// Paint grandparent (before rotate) to RED.
// [G] <P> [P]
// / \ rotate(P) / \ repaint / \
// <P> [U] ========> <X> [G] ======> <X> <G>
// / \ \
// <X> [U] [U]
rotate(p); //rotate x's parent
p->color = BLACK;
x->sibling()->color = RED; //repaint
删除操作
这个,不需要了! ——长崎素世
二叉查找树
我们先考虑单纯的二叉搜索树怎么删除一个节点。
根据要删除的节点的子节点数,可以分为三种情况:0个子节点(即叶节点),1个子节点(一条链的正中间),以及2个子节点(内部节点)。
删除叶节点是最简单的,直接把这个节点去掉就行了。
如果是1个子节点的情况,就与链表的删除比较类似,我们用这个子节点来代替原来被删除的节点。
2个子节点的情况,我们找到删除节点的后继节点,也就是右子树中权值最小的节点。找后继节点的方法是从右子节点开始一直往左走,直到没有左子节点为止。
(后继节点一定没有左子节点,也就是说,它只有0或1个子节点)
由于后继节点是原来右子树中最小的节点,所以把它原来的权值放到被删除节点的位置,仍然满足左子树 < 根 < 右子树的性质。
因此,我们交换删除节点和后继节点的权值,然后把后继节点删除,这就转化成了对有0或1个子节点的节点进行删除的情况。
红黑树
考虑删除有0或1个子节点的节点的情况。(有2个子节点的情况是可以转化成这两种情况的,所以不用讨论)
删除有1个子节点的节点时,是用它唯一的子节点代替本身。删除叶节点(0个子节点)时,可以看作用一个 NIL 节点代替它本身。
我们再考虑红黑树关键的两条性质:不能出现连续双红节点,以及黑高相同。
假如我们删除的是红节点,那么它的替代节点一定是黑节点,所以删除它既不会使树中出现双红节点,也不会影响黑高相同的性质。
但是如果删除的是黑节点,那么会使经过这个节点的路径的黑高-1,而且如果父节点、子节点都是红色,就会出现连续的双红了。
这里双红是很容易处理的,我们的重点在于处理黑高(“双黑”)。如果我们可以简单通过重新染色解决黑高不同的问题,那就简单处理就好了。但是有的时候这样行不通,我们就需要进行双黑修正。
具体地说,我们需要调整树的结构和颜色,把目标黑色节点放在一个比兄弟节点黑高多1的位置,再把目标节点删掉,这样黑高就相同了。
template<typename T>
bool RedBlackTree<T>::remove(T v, node* x) {
if(x == NULL) return false;
if(v != x->val) {
int branch = (v > x->val); //v > x->val : branch = 1, goto right child
if(x->ch[branch] != NULL) {
//the structure of the subtree may change
//node x may have new children after remove
//so first update the size of subtree
//if fail to remove then rollback size changes
x->siz--;
bool result = remove(v, x->ch[branch]);
if(result == false) {
x->siz++;
}
return result;
}
return false;
}
//Remove x from the tree
//......
}
Case 0
删除节点为根,且整棵树只有这一个节点,直接删就完了。
Case 1
删除节点 x x x 有2个子节点,则交换当前节点与后继节点的权值,然后问题转化为删除后继节点。
if(x->left != NULL && x->right != NULL) {
// Case 1: If the node is strictly internal
// Step 1. Find the successor S with the smallest key
// and its parent P on the right subtree.
// Step 2. Swap the data (key and value) of S and X,
// S is the node that will be deleted in place of X.
// Step 3. X = S, goto Case 2, 3
// | |
// X S
// / \ / \
// L .. swap(X, S) L ..
// | =========> |
// P P
// / \ / \
// S .. X ..
x->siz--;
//Step 1
node* rt = x->right;
rt->siz--;
while(rt->left) {
rt = rt->left;
rt->siz--;
}
//Step 2, 3
node* succ = rt;
swap(x->val, succ->val);
x = succ;
}
Case 2
删除叶子节点。如果是红叶子就直接删掉,如果是黑叶子,就要重新维护黑高。
在维护操作中,由于我们只是把目标节点放到了比兄弟节点黑高多1的地方,而没有改变目标节点的子树结构和数据,所以我们可以先对其维护,让它的黑高比兄弟多1,再把它删除。
if(x->left == NULL && x->right == NULL) {
// Case 2: Current node is a leaf
// Step 1. Put X to a position where its black height
// is greater than its sibling by 1.(if X is black)
// Step 2. remove X
// The maintain operation won't change the node itself,
// so we can perform maintain operation before unlink the node.
x->siz = 0;
if(x->color == BLACK) {
SolveDoubleBlack(x); //Step 1
}
x->father->ch[x->direction()] = NULL;
x->father->pushup();
return true;
}
Case 3
目标节点刚好有一个子节点,我们用它唯一的子节点来代替它本身。
这时唯一的孩子只可能是红色,否则一边为黑,一边为NIL,两边的黑高是不同的。
因此,我们直接把替代节点染成黑色,就可以解决黑高-1的问题。
// Case 3: Current node has a single left or right child
// Step 1. Paint its child to black(the child must be red).
// Step 2. Remove X
node* replacement = (x->left != NULL ? x->left : x->right);
if(x->color == BLACK) {
replacement->color = BLACK;
}
if(x == root) {
root = replacement;
replacement->father = NULL;
} else {
node* parent = x->father;
parent->ch[x->direction()] = replacement;
replacement->father = parent;
parent->pushup();
}
OI Wiki上这一段代码和最后给出的完整版里不一样,最后的完整版还写了唯一孩子是黑色时维护孩子的 if 分支,这个可能是出于工程上代码完整性的考虑(?
双黑修正
致命的黑影啊,翩翩起舞吧! ——Ave Mujica《黑色生日》
双黑维护,其实就是目标节点的黑高因为某些情况(如删除了一个黑节点)减少了1(或者即将要减少1),导致它比兄弟节点的黑高少1.
我们需要把它放在一个比兄弟黑高多1的位置,从而抵消黑高-1的影响。
Case 1
兄弟节点
s
s
s 为红色,则父节点
p
p
p 和两个侄节点
c
c
c 和
d
d
d 必为黑色。这种情况与双红修正的Case 5类似,无法直接使其满足所有性质,我们将其转化为其他Case进行处理。
Step 1. 往上转一次兄弟节点
s
s
s
Step 2. 把
s
s
s 染黑,
p
p
p 染红,并转到其他Case
假设原来黑高为上面最左边蓝色数字标注的值,则这样操作后目标节点与新的兄弟节点
c
c
c 黑高相同。由于我们需要目标节点比兄弟黑高多1,所以我们转到其它Case,继续对该节点进行处理。
if(sibling->color == RED) {
// Case 1: Sibling is RED, parent and nephews must be BLACK
// Step 1. Rotate X's sibling to P's position
// Step 2. Paint S to BLACK, P to RED
// Step 3. Goto Case 2, 3, 4, 5
// [P] <S> [S]
// / \ rotate(S) / \ repaint / \
// [X] <S> ==========> [P] [D] ======> <P> [D]
// / \ / \ / \
// [C] [D] [X] [C] [X] [C]
node* parent = x->father;
//Step 1
rotate(sibling);
//Step 2
sibling->color = BLACK;
parent->color = RED;
sibling = x->sibling(); //update sibling after rotation
}
Case 2
兄弟节点
s
s
s 和侄节点
c
c
c,
d
d
d 均为黑色,且父节点
p
p
p 为红色。
此时将父节点
p
p
p 染红,将
s
s
s 染黑即可。
假设原来的黑高如左图所示,那么重新染色之后,再删除目标节点,则其替代节点的黑高由2变为1。删除后,
p
p
p 的黑高为2。
假如
p
p
p 有一个父节点,则原来
p
p
p 的父节点黑高为3(因为原来
p
p
p 是红色),调整后父节点黑高也为3(
p
p
p 为黑,父节点黑高为
p
p
p 的黑高+1).
bool closeBlack = (closeNephew == NULL) || (closeNephew->color == BLACK);
bool distantBlack = (distantNephew == NULL) || (distantNephew->color == BLACK);
if(closeBlack && distantBlack) {
if(x->father->color == RED) {
// Case 2: Sibling and nephews are BLACK, parent is RED
// Swap the color of P and S
// <P> [P]
// / \ / \
// [X] [S] ====> [X] <S>
// / \ / \
// [C] [D] [C] [D]
sibling->color = RED;
x->father->color = BLACK;
}
//Other cases
}
Case 3
兄弟节点
s
s
s 和侄节点
c
c
c,
d
d
d 均为黑色,且父节点
p
p
p 也为黑色。
我们先把兄弟节点
s
s
s 染红,这样删除目标节点后,父节点
p
p
p 的所有黑高就相同了。但是这样
p
p
p 节点的黑高会比兄弟节点少1,所以我们递归维护
p
p
p。
重新染色后黑高的变化如上图所示。
//Assume that both nephews are black
if(x->father->color == BLACK) {
// Case 3: Sibling, parent and nephews are all black
// Step 1. Paint S to RED
// Step 2. Recursively maintain P
// [P] [P]
// / \ / \
// [X] [S] ====> [X] <S>
// / \ / \
// [C] [D] [C] [D]
sibling->color = RED;
SolveDoubleBlack(x->father);
return;
}
Case 4
我们把与目标节点同向(即同为左孩子/同为右孩子)的侄节点称为close nephew,用 c c c 表示。反向的侄节点称为distant nephew,用 d d d 表示。在图中可以看出,close nephew就是离目标节点(图中的N)比较近的侄节点,distant nephew就是离得比较远的侄节点。
Case 4为,兄弟节点
s
s
s 为黑色,
c
c
c 为红,
d
d
d 为黑。父节点
p
p
p 红黑均可。
这种情况同样无法直接维护,因此将其转变为 Case 5 的状态,利用后续 Case 5 的维护过程进行修正。
Step 1. 向上旋转
c
c
c
Step 2. 将
c
c
c 染黑,
s
s
s 染红。
Step 3. 转到 Case 5.
黑高变化如上图所示。
注:上面这个gif处理的是0003节点的双黑问题。
bool closeRed = (closeNephew != NULL) && (closeNephew->color == RED);
if(closeRed && distantBlack) {
// Case 4: Sibling is BLACK, close nephew is RED,
// distant nephew is BLACK
// Step 1. Rotate close nephew to sibling's position
// Step 2. Swap the color of close nephew and sibling
// Step 3. Goto case 5
// {P} {P}
// {P} / \ / \
// / \ rotate(C) [X] <C> repaint [X] [C]
// [X] [S] ==========> \ ======> \
// / \ [S] <S>
// <C> [D] \ \
// [D] [D]
//Step 1
rotate(closeNephew);
//Step 2
closeNephew->color = BLACK;
sibling->color = RED;
// Update sibling and nephews after rotation
sibling = x->sibling();
xdir = x->direction();
closeNephew = sibling->ch[xdir];
distantNephew = sibling->ch[xdir ^ 1];
}
Case 5
兄弟节点 s s s 为黑色,distant nephew节点 d d d 为红色,close nephew节点 c c c 为任意颜色,父节点 p p p 为任意颜色。
Step 1. 向上旋转兄弟节点
s
s
s
Step 2. 把
s
s
s 染成
p
p
p 的颜色,把
p
p
p 染黑(即交换两者颜色)。
Step 3. 将
d
d
d 染黑。
黑高变化如上图所示。这样操作后删除目标节点,目标节点黑高由2变为1,整棵树满足红黑树性质。
// Case 5: Sibling is BLACK, close nephew is unknown,
// distant nephew is RED
// {P} [S] {S}
// / \ rotate(S) / \ repaint / \
// [X] [S] ==========> {P} <D> =======> [P] [D]
// / \ / \ / \
// {C} <D> [X] {C} [X] {C}
// Step 1. Rotate sibling to P's position
// Step 2. Swap the color of parent and sibling.
// Paint distant nephew to BLACK if it is not null.
//Step 1
rotate(sibling);
//Step 2
sibling->color = x->father->color;
x->father->color = BLACK;
if(distantNephew != NULL) {
distantNephew->color = BLACK;
}
其他查询操作
其他查询操作就和二叉查找树完全一致了。下面几个查询操作对应洛谷P3369中的操作3-6.
查询排名
template<typename T>
int RedBlackTree<T>::get_rank(T v, node* x) {
if(x == NULL) return 0;
if(v <= x->val) return get_rank(v, x->left);
int lsiz = (x->left != NULL ? x->left->siz : 0);
return lsiz + 1 + get_rank(v, x->right);
}
template<typename T>
int RedBlackTree<T>::get_rank(T v) {
return get_rank(v, root);
}
查询第k大
template<typename T>
T RedBlackTree<T>::kth(int k, node* x) {
if(!(x->left)) {
if(k == 1) return x->val;
return kth(k - 1, x->right);
}
if(k <= x->left->siz) return kth(k, x->left);
if(k == x->left->siz + 1) return x->val;
return kth(k - x->left->siz - 1, x->right);
}
template<typename T>
T RedBlackTree<T>::kth(int k) {
return kth(k, root);
}
寻找前驱
template<typename T>
T RedBlackTree<T>::get_prev(T v) {
node *x = root;
T ans;
while(x != NULL) {
if(x->val < v) {
ans = x->val;
x = x->right;
} else {
x = x->left;
}
}
return ans;
}
寻找后继
template<typename T>
T RedBlackTree<T>::get_succ(T v) {
node *x = root;
T ans;
while(x != NULL) {
if(x->val > v) {
ans = x->val;
x = x->left;
} else {
x = x->right;
}
}
return ans;
}
最终代码
#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;
template<typename T>
class RedBlackTree {
private:
struct node;
node* root;
void SolveDoubleRed(node*);
void SolveDoubleBlack(node*);
node* Find(T);
void connect(node*, node*, int);
void rotate(node*);
void checkNodeSize(node*);
bool remove(T, node*);
int get_rank(T, node*);
T kth(int, node*);
public:
RedBlackTree() : root(NULL) {}
int size();
void insert(T);
bool remove(T);
int get_rank(T);
T kth(int);
T get_prev(T);
T get_succ(T);
void previs(node*);
void invis(node*);
void postvis(node*);
void print();
void checkNodeSize();
};
#define RED 1
#define BLACK 0
#define left ch[0]
#define right ch[1]
template<typename T>
struct RedBlackTree<T>::node {
/*
* <X> X is red.
* [X] X is black.
* {X} X is unknown(red/black).
*/
T val;
bool color; //1 is red, 0 is black
node *father, *ch[2];
int siz;
node(T v = T(), bool col = true, node* f = NULL,
node* l = NULL, node* r = NULL , int s = 1)
: val(v), color(col), father(f), siz(s) {
left = l;
right = r;
}
int direction() {
if(father == NULL) return 0;
return father->right == this;
}
node* sibling() {
if(father == NULL) return NULL;
return father->ch[direction() ^ 1];
}
bool hasSibling() {
return sibling() != NULL;
}
node* uncle() {
if(father == NULL) return NULL;
return father->sibling();
}
bool hasUncle() {
return uncle() != NULL;
}
void pushup() {
siz = (left?left->siz:0) + (right?right->siz:0) + 1;
}
void print() {
cout << "key = " << val << ", color = " << (color ? "Red" : "Black") << endl;
}
};
template<typename T>
int RedBlackTree<T>::size() {
return root->siz;
}
template<typename T>
void RedBlackTree<T>::connect(node *x, node *fa, int k) {
if(x != NULL) x->father = fa;
if(fa != NULL) {
fa->ch[k] = x;
} else {
root = x;
}
}
template<typename T>
void RedBlackTree<T>::rotate(node *x) {
//rotate x to its parent's position
node* y = x->father;
node* z = y->father;
int yson = x->direction();
if(z == NULL) {
root = x;
x->father = NULL;
} else {
int zson = y->direction();
connect(x, z, zson);
}
connect(x->ch[yson^1], y, yson);
connect(y, x, yson^1);
y->pushup();
x->pushup();
}
template<typename T>
void RedBlackTree<T>::insert(T v) {
node *x = root, *fa = NULL;
while(x != NULL) {
x->siz++;
fa = x;
if(v < x->val) {
x = x->left;
} else {
x = x->right;
}
}
x = new node(v, RED, fa); //create a new node
if(fa == NULL) {
root = x;
} else if(v < fa->val) {
fa->left = x;
} else {
fa->right = x;
}
SolveDoubleRed(x);
}
template<typename T>
void RedBlackTree<T>::SolveDoubleRed(node* x) {
if(x == root || x->father->color == BLACK) {
return;
}
node* p = x->father;
if(p == root) {
// Case 3: Parent is root and is RED
// Paint parent to BLACK.
// <P> [P]
// | ====> |
// <X> <X>
p->color = BLACK;
return;
}
if(x->hasUncle() && x->uncle()->color == RED) {
// Case 4: Both parent and uncle are RED
// Paint parent and uncle to BLACK;
// Paint grandparent to RED;
// Maintain grandparent recursively.
// [G] <G>
// / \ / \
// <P> <U> ====> [P] [U]
// / /
// <X> <X>
p->color = BLACK; //parent -> black
x->uncle()->color = BLACK; //uncle -> black
p->father->color = RED; //grandparent -> red
SolveDoubleRed(p->father);
return;
}
// Case 5 & 6: parent is RED, uncle is BLACK(or NULL)
if(x->direction() != p->direction()) {
// Case 5: Current node is the opposite direction as parent
// Step 1. Rotate x to parent's position.
// Step 2. Goto Case 6.
// [G] [G]
// / \ rotate(X) / \
// <P> [U] ========> <X> [U]
// \ /
// <X> <P>
rotate(x);
x = p; //Now P is the child of double red.
p = x->father; //reset p to x's father
}
// Case 6: Current node is the same direction as parent
// Step 1. Rotate parent to grandparent's position
// Step 2. Paint parent (before rotate) to BLACK;
// Paint grandparent (before rotate) to RED.
// [G] <P> [P]
// / \ rotate(P) / \ repaint / \
// <P> [U] ========> <X> [G] ======> <X> <G>
// / \ \
// <X> [U] [U]
rotate(p); //rotate x's parent
p->color = BLACK;
x->sibling()->color = RED; //repaint
}
#define col(a) (a == RED ? "Red" : "Black")
template<typename T>
void RedBlackTree<T>::previs(node* x) {
if(x == NULL) return;
printf("%d %s %d\n", x->val, col(x->color), x->siz);
previs(x->left);
previs(x->right);
}
template<typename T>
void RedBlackTree<T>::invis(node* x) {
if(x == NULL) return;
invis(x->left);
printf("%d %s %d\n", x->val, col(x->color), x->siz);
invis(x->right);
}
template<typename T>
void RedBlackTree<T>::postvis(node* x) {
if(x == NULL) return;
postvis(x->left);
postvis(x->right);
printf("%d %s %d\n", x->val, col(x->color), x->siz);
}
template<typename T>
void RedBlackTree<T>::print() {
printf("------pre-vis------\n");
previs(root);
printf("------in-vis------\n");
invis(root);
printf("------post-vis------\n");
postvis(root);
}
template<typename T>
void RedBlackTree<T>::checkNodeSize(node* x) {
int before = x->siz;
if(x->left) checkNodeSize(x->left);
if(x->right) checkNodeSize(x->right);
x->pushup();
if(x->siz != before) {
printf("node of key %d : size changed from %d to %d\n", x->val, before, x->siz);
}
}
template<typename T>
void RedBlackTree<T>::checkNodeSize() {
checkNodeSize(root);
}
template<typename T>
bool RedBlackTree<T>::remove(T v) {
return remove(v, root);
}
template<typename T>
bool RedBlackTree<T>::remove(T v, node* x) {
if(x == NULL) return false;
if(v != x->val) {
int branch = (v > x->val); //v > x->val : branch = 1, goto right child
if(x->ch[branch] != NULL) {
//the structure of the subtree may change
//node x may have new children after remove
//so first update the size of subtree
//if fail to remove then rollback size changes
x->siz--;
bool result = remove(v, x->ch[branch]);
if(result == false) {
x->siz++;
}
return result;
}
return false;
}
if(x == root && x->siz == 1) {
root = NULL;
return true;
}
if(x->left != NULL && x->right != NULL) {
// Case 1: If the node is strictly internal
// Step 1. Find the successor S with the smallest key
// and its parent P on the right subtree.
// Step 2. Swap the data (key and value) of S and X,
// S is the node that will be deleted in place of X.
// Step 3. X = S, goto Case 2, 3
// | |
// X S
// / \ / \
// L .. swap(X, S) L ..
// | =========> |
// P P
// / \ / \
// S .. X ..
x->siz--;
//Step 1
node* rt = x->right;
rt->siz--;
while(rt->left) {
rt = rt->left;
rt->siz--;
}
//Step 2, 3
node* succ = rt;
swap(x->val, succ->val);
x = succ;
}
if(x->left == NULL && x->right == NULL) {
// Case 2: Current node is a leaf
// Step 1. Put X to a position where its black height
// is greater than its sibling by 1.(if X is black)
// Step 2. remove X
// The maintain operation won't change the node itself,
// so we can perform maintain operation before unlink the node.
x->siz = 0;
if(x->color == BLACK) {
SolveDoubleBlack(x); //Step 1
}
x->father->ch[x->direction()] = NULL;
x->father->pushup();
return true;
}
// Case 3: Current node has a single left or right child
// Step 1. Paint its child to black(the child must be red).
// Step 2. Remove X
node* replacement = (x->left != NULL ? x->left : x->right);
if(x->color == BLACK) {
replacement->color = BLACK;
}
if(x == root) {
root = replacement;
replacement->father = NULL;
} else {
node* parent = x->father;
parent->ch[x->direction()] = replacement;
replacement->father = parent;
parent->pushup();
}
return true;
}
template<typename T>
void RedBlackTree<T>::SolveDoubleBlack(node* x) {
if(x == root) return;
node* sibling = x->sibling();
if(sibling->color == RED) {
// Case 1: Sibling is RED, parent and nephews must be BLACK
// Step 1. Rotate X's sibling to P's position
// Step 2. Paint S to BLACK, P to RED
// Step 3. Goto Case 2, 3, 4, 5
// [P] <S> [S]
// / \ rotate(S) / \ repaint / \
// [X] <S> ==========> [P] [D] ======> <P> [D]
// / \ / \ / \
// [C] [D] [X] [C] [X] [C]
node* parent = x->father;
//Step 1
rotate(sibling);
//Step 2
sibling->color = BLACK;
parent->color = RED;
sibling = x->sibling(); //update sibling after rotation
}
//close nephew: sibling's child with the same direction as x
int xdir = x->direction(); //the direction of x
node* closeNephew = sibling->ch[xdir];
node* distantNephew = sibling->ch[xdir ^ 1];
//NIL nodes are always black
bool closeBlack = (closeNephew == NULL) || (closeNephew->color == BLACK);
bool distantBlack = (distantNephew == NULL) || (distantNephew->color == BLACK);
if(closeBlack && distantBlack) {
if(x->father->color == RED) {
// Case 2: Sibling and nephews are BLACK, parent is RED
// Swap the color of P and S
// <P> [P]
// / \ / \
// [X] [S] ====> [X] <S>
// / \ / \
// [C] [D] [C] [D]
sibling->color = RED;
x->father->color = BLACK;
} else {
// Case 3: Sibling, parent and nephews are all black
// Step 1. Paint S to RED
// Step 2. Recursively maintain P
// [P] [P]
// / \ / \
// [X] [S] ====> [X] <S>
// / \ / \
// [C] [D] [C] [D]
sibling->color = RED;
SolveDoubleBlack(x->father);
}
} else {
bool closeRed = (closeNephew != NULL) && (closeNephew->color == RED);
if(closeRed && distantBlack) {
// Case 4: Sibling is BLACK, close nephew is RED,
// distant nephew is BLACK
// Step 1. Rotate close nephew to sibling's position
// Step 2. Swap the color of close nephew and sibling
// Step 3. Goto case 5
// {P} {P}
// {P} / \ / \
// / \ rotate(C) [X] <C> repaint [X] [C]
// [X] [S] ==========> \ ======> \
// / \ [S] <S>
// <C> [D] \ \
// [D] [D]
//Step 1
rotate(closeNephew);
//Step 2
closeNephew->color = BLACK;
sibling->color = RED;
// Update sibling and nephews after rotation
sibling = x->sibling();
xdir = x->direction();
closeNephew = sibling->ch[xdir];
distantNephew = sibling->ch[xdir ^ 1];
}
// Case 5: Sibling is BLACK, close nephew is unknown,
// distant nephew is RED
// {P} [S] {S}
// / \ rotate(S) / \ repaint / \
// [X] [S] ==========> {P} <D> =======> [P] [D]
// / \ / \ / \
// {C} <D> [X] {C} [X] {C}
// Step 1. Rotate sibling to P's position
// Step 2. Swap the color of parent and sibling.
// Paint distant nephew to BLACK if it is not null.
//Step 1
rotate(sibling);
//Step 2
sibling->color = x->father->color;
x->father->color = BLACK;
if(distantNephew != NULL) {
distantNephew->color = BLACK;
}
}
}
template<typename T>
T RedBlackTree<T>::kth(int k, node* x) {
if(!(x->left)) {
if(k == 1) return x->val;
return kth(k - 1, x->right);
}
if(k <= x->left->siz) return kth(k, x->left);
if(k == x->left->siz + 1) return x->val;
return kth(k - x->left->siz - 1, x->right);
}
template<typename T>
T RedBlackTree<T>::kth(int k) {
return kth(k, root);
}
template<typename T>
T RedBlackTree<T>::get_prev(T v) {
node *x = root;
T ans;
while(x != NULL) {
if(x->val < v) {
ans = x->val;
x = x->right;
} else {
x = x->left;
}
}
return ans;
}
template<typename T>
T RedBlackTree<T>::get_succ(T v) {
node *x = root;
T ans;
while(x != NULL) {
if(x->val > v) {
ans = x->val;
x = x->left;
} else {
x = x->right;
}
}
return ans;
}
template<typename T>
int RedBlackTree<T>::get_rank(T v, node* x) {
if(x == NULL) return 0;
if(v <= x->val) return get_rank(v, x->left);
int lsiz = (x->left != NULL ? x->left->siz : 0);
return lsiz + 1 + get_rank(v, x->right);
}
template<typename T>
int RedBlackTree<T>::get_rank(T v) {
return get_rank(v, root);
}
RedBlackTree<int> Tree;
int main() {
int n;
cin >> n;
while(n--) {
int op, x;
cin >> op >> x;
if(op == 1) {
Tree.insert(x);
} else if(op == 2) {
Tree.remove(x);
} else if(op == 3) {
cout << Tree.get_rank(x) + 1 << endl;
} else if(op == 4) {
cout << Tree.kth(x) << endl;
} else if(op == 5) {
cout << Tree.get_prev(x) << endl;
} else {
cout << Tree.get_succ(x) << endl;
}
// Tree.checkNodeSize();
}
return 0;
}