红黑树的深入理解与实现 | 递归插入及删除
红黑树在Linux，Java和C++中都被作为数据结构的底层实现，一直以来都是面试吐槽点，手撕红黑树的段子广为流传，这两天我学习了一下
Algorithms 4th Edition中的左斜红黑树
该书的作者Robert Sedgewick恰恰也是左斜红黑树的发明者
定义
首先，红黑树是一种二叉查找树，基于2-3树的高效平衡插入思想实现，同时又兼具二叉查找树简洁的查找方法，下图中分别为2-3树和对应的红黑树
2-3树中包括2-结点和3-结点，表示结点最多有几个孩子
红黑树是指用红色的链连接一个3-结点，不过为了便于实现，将链的颜色保存在结点中
2-3树和红黑树之间可以相互转换，红黑树中的红色边表示一个3-结点，将所有红色结点进行提升变换：
通过提升变换就将红黑树转换成2-3树的形式，每个结点到叶子的黑色高度都是相等的。在左斜红黑树中规定红色边必须是左斜的，我们需要保证无论插入还是删除都不能打破这个规则
根据上面的转换，也可以归纳出红黑树的性质如下：
- 每个结点只能是红色或黑色
- 根结点一定为黑色
- 所有叶子结点为黑色空结点
- 每个红色结点的父结点和子结点必为黑色
- 从任一结点到其每个叶子的路径上黑色结点数相同
依此可以定义红黑树结构如下：
type color bool
const (
Red = true
Black = false
)
// Comparator 用于比较key
type Comparator interface {
CompareTo(Comparator) int
}
type node struct {
key Comparator
val interface{}
left *node
right *node
color color
size int
}
// NewNode ...
func NewNode(key Comparator, val interface{}, color color, size int) *node {
return &node{
key: key,
val: val,
color: color,
size: size,
}
}
// RedBlackBST 红黑树
type RedBlackBST struct {
root *node
}
// NewRedBlackBST ...
func NewRedBlackBST() *RedBlackBST {
return new(RedBlackBST)
}
func (t *RedBlackBST) isEmpty() bool {
return t.root == nil
}
// Size 红黑树结点数量
func (t *RedBlackBST) Size() int {
return t.root.Size()
}
// Size n为根的子树结点数量
func (n *node) Size() int {
if n == nil {
return 0
}
return n.size
}
func (n *node) isRed() bool {
if n == nil {
return false
}
return n.color == Red
}
var ErrUnderflow = errors.New("BST underflow")
var ErrArgumentNil = errors.New("argument to delete is nil")
查找
根据定义可知红黑树的查找和普通的二叉查找树是一样的
// Get 返回key对应的val
func (t *RedBlackBST) Get(key Comparator) (interface{}, error) {
if key == nil {
return nil, ErrArgumentNil
}
return t.root.get(key), nil
}
func (n *node) get(key Comparator) interface{} {
for n != nil {
if cmp := key.CompareTo(n.key); cmp < 0 {
n = n.left
} else if cmp > 0 {
n = n.right
} else {
return n.val
}
}
return nil
}
插入
首先要考虑的是插入结点的颜色
如果插入黑色结点，从该结点到根会比其他路径上的黑色结点数多一个，打破了第5条规则，需要对整颗树进行平衡调整
如果插入红色结点，只需要考虑如何修复第4条规则，可以将调整限制在很小的子树范围内
所以将新插入的结点置为红色
然后需要进行平衡调整
可以根据对应2-3树中待插入父节点的类型分为两种情况
父节点为2-结点
这里又分为两种情况，插入位置在左边和右边
- 如果插入到左边，恰好保证了左斜，无需调整
-
如果插入到右边，则需要通过左旋将其从右斜变为左斜，并互换颜色，定义这一过程为rotateLeft
父节点为3-结点
因为有三条子边，插入位置可以分为左边、中间和右边三种情况
- 右边
插入结点E，逆置三个结点颜色，相当于提升3-结点中间元素到上一层，定义为flipColors -
左边
先将B和D右旋，互换颜色，定义为rotateRight，再逆置颜色flipColors -
中间
先将B和C进行rotateLeft，再将C和D进行rotateRight，最后flipColors
其实在递归实现的时候不用管最后这种情况，因为B和C进行rotateLeft就相当于2阶结点的第二种情况，之后按3阶左边情况处理即可（也就是相当于B的2阶右边+D的3阶左边的组合）
三种操作演示
完整流程
蓝色箭头表示插入结点
代码实现
// Put 插入key与val
func (t *RedBlackBST) Put(key Comparator, val interface{}) error {
if key == nil {
return ErrArgumentNil
}
// val为空则删除
if val == nil {
return t.Delete(key)
}
t.root = t.root.put(key, val)
// 保证root为黑色
t.root.color = Black
return nil
}
func (n *node) put(key Comparator, val interface{}) *node {
if n == nil {
return NewNode(key, val, Red, 1)
}
// 插入元素
if cmp := key.CompareTo(n.key); cmp < 0 {
n.left = n.left.put(key, val)
} else if cmp > 0 {
n.right = n.right.put(key, val)
} else {
n.val = val
}
// 平衡调整
if !n.left.isRed() && n.right.isRed() {
n = n.rotateLeft()
}
if n.left.isRed() && n.left.left.isRed() {
n = n.rotateRight()
}
if n.left.isRed() && n.right.isRed() {
n.flipColors()
}
n.size = n.left.Size() + n.right.Size() + 1
return n
}
// 右旋 互换颜色
func (n *node) rotateRight() *node {
x := n.left
x.right, n.left = n, x.right
x.color, n.color = n.color, x.color
x.size, n.size = n.Size(), n.left.Size() + n.right.Size() + 1
return x
}
// 左旋 互换颜色
func (n *node) rotateLeft() *node {
x := n.right
x.left, n.right = n, x.left
x.color, n.color = n.color, x.color
x.size, n.size = n.Size(), n.left.Size() + n.right.Size() + 1
return x
}
// 逆置颜色
func (n *node) flipColors() {
n.color, n.left.color, n.right.color = !n.color, !n.left.color, !n.right.color
}
删除
删除相对来说更为复杂，书中课后题代码写的就有问题，而很多博客写的更是错上加错，浪费了我几个小时也没想明白，最后找到官网源码才真正弄明白
deleteMin
首先考虑如何删除最小值点，有两种情况，如果待删除点是3-结点可以直接删除变成2-结点，然后做局部调整即可；如果待删除点是2-结点，删除会导致整个树的结构发生变换，这种情况是难以处理的，所以需要将2-结点转换为3-结点
这时候需要“借”它的父结点下去临时组成4-结点，也就是插入时提升元素的逆过程，同样可以用flipColors表示
记父节点为n，如果n.right.left是红色结点，说明n.right之前已经是3-结点，此时组成的就是5-结点了，需要进行如下变换把n.right提升到上一层
以上步骤将父节点或兄弟结点的红色结点借给左子树，因此可以定义为moveRedLeft，与之相反的操作为moveRedRight
当!n.left.isRed() && !n.left.left.isRed()，也就是n.left为2-结点时，调用moveRedLeft，然后就可以用递归的方式删除min结点，删除之后还要重新进行局部平衡调整（类似插入的时候）
deleteMax
deleteMax与deleteMin的处理类似，不过由于红链接都是左斜的，也就意味着3-结点中的较大者与父节点直接相连，不能直接删除，需要先rotateRight临时转换成右斜红链接，删除之后再调整回来
delete
如果待删除的key小于当前结点的key，就相当于进行了deleteMin操作，否则相当于deleteMax
当key相等时，交换当前结点n和n右子树中的最小结点，然后删除min结点，因为min结点一定是叶子结点，删除造成的影响较小，详细思路见注释
// DeleteMin 删除最小值
func (t *RedBlackBST) DeleteMin() error {
if t.isEmpty() {
return ErrUnderflow
}
if !t.root.left.isRed() && !t.root.right.isRed() {
t.root.color = Red
}
t.root = t.root.deleteMin()
if !t.isEmpty() {
t.root.color = Black
}
return nil
}
func (n *node) deleteMin() *node {
// 没有左子树说明n最小 直接返回空
if n.left == nil {
return nil
}
// n.left为2-结点 需要借一个红色结点变成3-结点
if !n.left.isRed() && !n.left.left.isRed() {
n = n.moveRedLeft()
}
n.left = n.left.deleteMin()
return n.balance()
}
// DeleteMax 删除最大值
func (t *RedBlackBST) DeleteMax() error {
if t.isEmpty() {
return ErrUnderflow
}
if !t.root.left.isRed() && !t.root.right.isRed() {
t.root.color = Red
}
t.root = t.root.deleteMax()
if !t.isEmpty() {
t.root.color = Black
}
}
func (n *node) deleteMax() *node {
// 左斜红链接转为右斜 父节点指向3-结点较小者
if n.left.isRed() {
n = n.rotateRight()
}
// 没有右子树说明n最大 直接返回空
if n.right == nil {
return nil
}
// n.right为2-结点 需要借一个红色结点变成3-结点
if !n.right.isRed() && !n.right.left.isRed() {
n = n.moveRedRight()
}
n.right = n.right.deleteMax()
return n.balance()
}
// Delete 删除key对应的结点
func (t *RedBlackBST) Delete(key Comparator) error {
if key == nil {
return ErrArgumentNil
}
// 找不到直接退出
if t.root.get(key) == nil {
return nil
}
if !t.root.left.isRed() && !t.root.right.isRed() {
t.root.color = Red
}
t.root = t.root.delete(key)
if !t.isEmpty() {
t.root.color = Black
}
return nil
}
func (n *node) delete(key Comparator) *node {
// 如果key比当前结点的key小
if key.CompareTo(n.key) < 0 {
// 借红色结点以保证左子树有3-结点
if !n.left.isRed() && !n.left.left.isRed() {
n = n.moveRedLeft()
}
// 递归调用n.left
n.left = n.left.delete(key)
} else {
// 将父节点指向3-结点中的较小值 便于删除较大值
if n.left.isRed() {
n = n.rotateRight()
}
// 如果找到目标结点且没有比它更大的结点 直接返回空表示删除
if key.CompareTo(n.key) == 0 && n.right == nil {
return nil
}
// 借红色结点以保证右子树有3-结点
if !n.right.isRed() && !n.right.left.isRed() {
n = n.moveRedRight()
}
// 如果找到目标且有比他更大的结点
// 交换n与其相邻的较大值 也就是n.right.min()
// 然后调用deleteMin删除min结点即可
if key.CompareTo(n.key) == 0 {
x := n.right.min()
n.key, n.val = x.key, x.val
n.right = n.right.deleteMin()
} else {
// 否则继续递归调用n.right
n.right = n.right.delete(key)
}
}
return n.balance()
}
func (n *node) moveRedLeft() *node {
n.flipColors()
if n.right.left.isRed() {
n.right = n.right.rotateRight()
n = n.rotateLeft()
n.flipColors()
}
return n
}
func (n *node) moveRedRight() *node {
n.flipColors()
if n.left.left.isRed() {
n = n.rotateRight()
n.flipColors()
}
return n
}
// 平衡调整
func (n *node) balance() *node {
if n.right.isRed() {
n = n.rotateLeft()
}
if n.left.isRed() && n.left.left.isRed() {
n = n.rotateRight()
}
if n.left.isRed() && n.right.isRed() {
n.flipColors()
}
n.size = n.left.Size() + n.right.Size() + 1
return n
}
// 返回子树中最小结点
func (n *node) min() *node {
if n.left == nil {
return n
}
return n.left.min()
}
测试
package rbt
import (
"math/rand"
"testing"
"time"
)
var (
rng = rand.New(rand.NewSource(time.Now().Unix()))
tree = NewRedBlackBST()
N int = 1e3
min int = 0
max int = 1e5
)
type key int
func (v key) CompareTo(i Comparator) int {
if v < i.(key) {
return -1
}
if v > i.(key) {
return 1
}
return 0
}
func randKey() Comparator {
return key(rng.Intn(max-min-1) + min + 1)
}
func TestRedBlackBST(t *testing.T) {
m := make(map[Comparator]interface{})
// put
for i := 0; i < N; i++ {
k, v := randKey(), i
_ = tree.Put(k, v)
m[k] = v
verify(t)
}
// get
for i := 0; i < N; i++ {
k := randKey()
v, _ := tree.Get(k)
compareVal(t, v, m[k])
}
// delete
for i := 0; i < N; i++ {
k := randKey()
_ = tree.Delete(k)
verify(t)
}
}
func compareVal(t *testing.T, v1, v2 interface{}) {
if v1 != v2 {
t.Error("The value found does not match")
}
}
func verify(t *testing.T) {
if !tree.isSizeConsistent() {
t.Error("Subtree counts not consistent")
}
// test BST
if !tree.isBST() {
t.Error("Not in symmetric order")
}
// test RBT
if !tree.isRootBlack() {
t.Error("Root is not black")
}
if !tree.is23() {
t.Error("Not a 2-3 tree")
}
if !tree.isBalanced() {
t.Error("Not balanced")
}
}
func (t *RedBlackBST) isBST() bool {
return t.root.isBST(key(min), key(max))
}
func (n *node) isBST(min, max Comparator) bool {
if n == nil {
return true
}
if n.key.CompareTo(min) <= 0 {
return false
}
if n.key.CompareTo(max) >= 0 {
return false
}
return n.left.isBST(min, n.key) && n.right.isBST(n.key, max)
}
func (t *RedBlackBST) isSizeConsistent() bool {
return t.root.isSizeConsistent()
}
func (n *node) isSizeConsistent() bool {
if n == nil {
return true
}
if n.size != n.left.Size()+n.right.Size()+1 {
return false
}
return n.left.isSizeConsistent() && n.right.isSizeConsistent()
}
func (t *RedBlackBST) isRootBlack() bool {
return !t.root.isRed()
}
func (t *RedBlackBST) is23() bool {
return t.root.is23()
}
func (n *node) is23() bool {
if n == nil {
return true
}
if n.right.isRed() {
return false
}
if n.isRed() && n.left.isRed() {
return false
}
return n.left.is23() && n.right.is23()
}
func (t *RedBlackBST) isBalanced() bool {
black := 0
x := t.root
for x != nil {
if !x.isRed() {
black++
}
x = x.left
}
return t.root.isBalanced(black)
}
func (n *node) isBalanced(black int) bool {
if n == nil {
return black == 0
}
if !n.isRed() {
black--
}
return n.left.isBalanced(black) && n.right.isBalanced(black)
}
