红黑树代码实现及注释

代码参考了以下两个链接：

• https://www.programiz.com/dsa/deletion-from-a-red-black-tree
• https://github.com/openjdk/jdk/blob/master/src/java.base/share/classes/java/util/TreeMap.java

主要是将 TreeMap 的实现转化为 Go 代码，并且为了容易理解，把 TreeMap 的代码转为不优雅的形式。

理解过程中借助了图形展示：

https://www.cs.usfca.edu/~galles/visualization/RedBlack.html

不确保理解完全正确。

完整代码及注释

package redblacktree

import "math"

const (
    COLOR_BLACK = 0
    COLOR_RED   = 1
)

type TreeNode struct {
    Value int
    Color int

    Parent *TreeNode
    Left   *TreeNode
    Right  *TreeNode
}

// NULL_NODE 作为通用叶子节点使代码在获取 TreeNode 属性前无需判断是否为 nil
var NULL_NODE *TreeNode

func NewTreeNode(value int, parent *TreeNode) *TreeNode {
    return &TreeNode{
        Value:  value,
        Color:  COLOR_RED,
        Parent: parent,
        Left:   NULL_NODE,
        Right:  NULL_NODE,
    }
}

type RedBlackTree struct {
    root *TreeNode
}

func NewRedBlackTree() *RedBlackTree {
    NULL_NODE = &TreeNode{}
    NULL_NODE.Color = COLOR_BLACK
    NULL_NODE.Value = math.MaxInt32
    NULL_NODE.Parent = NULL_NODE
    NULL_NODE.Left = NULL_NODE
    NULL_NODE.Right = NULL_NODE

    return &RedBlackTree{root: NULL_NODE}
}

func (tree *RedBlackTree) Insert(value int) {
    if tree.root == nil {
        tree.root = NewTreeNode(value, nil)
        tree.root.Color = COLOR_BLACK
        return
    }

    var parent *TreeNode

    current := tree.root
    for current != NULL_NODE {
        // 这里重点在寻找 parent
        parent = current

        if value < current.Value {
            current = current.Left
        } else if value > current.Value {
            current = current.Right
        } else {
            return
        }
    }

    newNode := NewTreeNode(value, parent)
    if value < parent.Value {
        parent.Left = newNode
    } else {
        parent.Right = newNode
    }

    if parent.Parent == nil {
        return
    }

    tree.fixAfterInsertion(newNode)
}

func (tree *RedBlackTree) fixAfterInsertion(newNode *TreeNode) {
    current := newNode // Color == RED
    if current.Parent.Color == COLOR_BLACK {
        return
    }

    // 循环主要是因为翻转颜色的场景时，current 更新为 grandparent
    for current.Parent.Color == COLOR_RED {
        parent := current.Parent     // Color == RED
        grandParent := parent.Parent // Color == BLACK
        var uncle *TreeNode

        if parent == grandParent.Left {
            //        grandParent(BALCK)
            //          /           \
            //    parent(RED)       uncle(?)

            uncle = grandParent.Right

            if uncle.Color == COLOR_RED { // 祖父一黑带两红，翻转颜色
                //        grandParent(BALCK)
                //          /           \
                //    parent(RED)       uncle(RED)

                uncle.Color = COLOR_BLACK
                parent.Color = COLOR_BLACK
                grandParent.Color = COLOR_RED
                //        grandParent(RED)
                //          /           \
                //    parent(BALCK)       uncle(BALCK)

                current = grandParent
            } else if uncle.Color == COLOR_BLACK {
                // 翻转颜色后出现的情况，祖父左红右黑，需要旋转

                // 红色父子不同边，先左旋
                if current == parent.Right {
                    //         grandParent(BALCK)
                    //          /           \
                    //       parent(RED)       uncle(BALCK)
                    //         /      \
                    // sibling(BLACK)  current(RED)

                    RotateLeft(parent)
                    //        grandParent(BALCK)
                    //          /           \
                    //    current(RED)      uncle(BALCK)
                    //        /        \
                    //     parent(RED)   right(BLACK)
                    //    /           \
                    // sibling(BLACK)  left(BLACK)

                    current = parent
                    parent = current.Parent
                    grandParent = parent.Parent // 和之前同一个节点
                    //        grandParent(BALCK)
                    //          /           \
                    //    parent(RED)      uncle(BALCK)
                    //     /        \
                    // current(RED) sibling(BLACK)
                }

                //        grandParent(BALCK)
                //          /           \
                //    parent(RED)      uncle(BALCK)
                //     /        \
                // current(RED)  sibling(BLACK)

                current.Parent.Color = COLOR_BLACK
                grandParent.Color = COLOR_RED
                //        grandParent(RED)
                //          /           \
                //    parent(BALCK)      uncle(BALCK)
                //     /        \
                // current(RED)  sibling(BLACK)

                subRoot := RotateRight(grandParent)
                //              parent(BALCK)
                //          /                   \
                //   current(RED)                 grandParent(RED)
                //     /        \                  /            \
                // left(BLACK) right(BLACK)   sibling(BLACK)     uncle(BALCK)
                if subRoot.Parent == nil {
                    tree.root = subRoot
                }
            }
        } else if parent == grandParent.Right {
            //        grandParent(BALCK)
            //          /           \
            //    uncle(?)       parent(RED)

            uncle = grandParent.Left
            if uncle.Color == COLOR_RED {
                parent.Color = COLOR_BLACK
                uncle.Color = COLOR_BLACK
                grandParent.Color = COLOR_RED

                current = grandParent
            } else {
                if current == parent.Left {
                    RotateRight(parent)

                    current = parent
                    parent = current.Parent
                    grandParent = parent.Parent
                }

                parent.Color = COLOR_BLACK
                grandParent.Color = COLOR_RED
                subRoot := RotateLeft(grandParent)
                if subRoot.Parent == nil {
                    tree.root = subRoot
                }
            }
        }

        // 仅翻转时出现。current 更新为 grandparent，可能为根节点或者其父节点是根节点。
        // 如果其父节点是根节点，由于根节点必为黑，会跳出循环，不必再加判断其祖父节点是否存在。
        if current == tree.root {
            break
        }
    }

    // 无论如何旋转或者翻转，都要设置根节点为黑色
    tree.root.Color = COLOR_BLACK
}

func (tree *RedBlackTree) Get(value int) *TreeNode {
    if tree.root == nil {
        return nil
    }

    current := tree.root
    for current != NULL_NODE {
        if value < current.Value {
            current = current.Left
        } else if value > current.Value {
            current = current.Right
        } else {
            return current
        }
    }

    return nil
}

func (tree *RedBlackTree) Delete(value int) {
    target := tree.Get(value)
    if target == nil {
        return
    }

    // 情况1：target 有两个子节点。取中序遍历下一个节点值覆盖当前节点值，交换两个节点的值
    if target.Left != NULL_NODE && target.Right != NULL_NODE {
        successor := tree.Successor(target)
        target.Value = successor.Value

        // 要删除的节点变更为刚刚找到的位置
        // 由于即将被删除，因此就算不严格地执行 value 交换也没关系
        target = successor
    }

    // 情况2：target 至少一个子节点为空。再分为两种情况：其中一个不为空；两个都为空。
    // 如果之前满足情况 1，节点交换后 target 就会变成情况 2。
    var replacement *TreeNode
    if target.Left == NULL_NODE {
        replacement = target.Right
    } else {
        replacement = target.Left
    }

    // 情况3：target 其中一个子节点为空，另一个子节点（即 replacement）不为空，且必为红
    // 由于子节点有红色节点，则 target 必为黑。
    // 让 replacement 替换 target，此时 replacement 仍然为红，删除 target。
    if replacement != NULL_NODE {
        replacement.Parent = target.Parent
        if target.Parent == nil {
            tree.root = replacement
        } else if target == target.Parent.Left {
            target.Parent.Left = replacement
        } else {
            target.Parent.Right = replacement
        }

        target.Left = nil
        target.Right = nil
        target.Parent = nil

        // 修复红色节点
        tree.fixAfterDeletion(replacement)
        return
    }

    // 两个子节点都为 nil，分为两种情况：

    // 1. 如果是根节点，则清空
    if target.Parent == nil {
        tree.root = nil
        return
    }

    // 2. 不是根节点。如果是黑色节点，会破坏红黑树性质，需要修复。
    // 这里假设用它本身作为替代（上面的 replacement），后面再删除它
    if target.Color == COLOR_BLACK {
        tree.fixAfterDeletion(target)
    }

    // 如果无需修复，说明就是红色。红色节点可以直接删除。
    // 如果需要修复，修复后 target 为单边黑色子节点，可以直接删除。
    if target.Parent != nil {
        if target == target.Parent.Left {
            target.Parent.Left = NULL_NODE
        } else if target == target.Parent.Right {
            target.Parent.Right = NULL_NODE
        }

        target.Parent = nil
    }
}

func (tree *RedBlackTree) fixAfterDeletion(node *TreeNode) {
    // node 的颜色可为黑或红
    if node.Color == COLOR_RED {
        node.Color = COLOR_BLACK
        return
    }

    var sibling *TreeNode
    for node != tree.root && node.Color == COLOR_BLACK {
        parent := node.Parent // 无视 parent 当前的颜色

        /** 接下来分别处理 node 在 parent 左右两边的情况 */

        /** 情况一：node 在 parent 左边 */
        if node == parent.Left {
            // 如果一个非空节点是黑的，它必然存在非空兄弟
            sibling = parent.Right

            // 如果黑色节点的兄弟是红色，那么：
            // 1. 兄弟必然有两个黑色儿子
            // 2. 父节点必然是黑色
            if sibling.Color == COLOR_RED {
                //          parent(BLACK)
                //          /           \
                //  node(BLACK)         sibling(RED)
                //     /    \           /           \
                //   (?)    (?) sibling_left(BLACK) sibling_right(BLACK)
                //                  /     \             /    \
                //                (?)     (?)         (?)    (?)

                sibling.Color = COLOR_BLACK
                parent.Color = COLOR_RED
                //            parent(RED)
                //          /           \
                //  node(BLACK)         sibling(BLACK)
                //     /    \           /           \
                //   (?)    (?) sibling_left(BLACK) sibling_right(BLACK)

                RotateLeft(parent)
                //          sibling(BLACK)
                //          /           \
                //    parent(RED)       sibling_right(BLACK)
                //     /       \
                // node(BLACK)  sibling_left(BLACK)

                sibling = parent.Right
                //          grandparent(BLACK)
                //          /           \
                //    parent(RED)       uncle(BLACK)
                //     /       \
                // node(BLACK)  sibling(BLACK)
            }

            /** 到这里时，sibling 节点必然为黑色，但可能有红色子节点 */

            // 由于 node 黑色，减少了路径的黑色节点数量。可以通过两种方式调整父节点左右两边黑色节点数：
            // 1. 通过旋转补充黑色节点
            // 2. 通过让兄弟子树某个黑节点变红减少黑色节点

            // 兄弟左右子节点都为黑，包含了 NULL_NODE 的情况
            // 把黑色节点往上提，减少兄弟侧的黑色节点数
            if sibling.Left.Color == COLOR_BLACK && sibling.Right.Color == COLOR_BLACK {
                sibling.Color = COLOR_RED

                // 如果 sibling 一开始是红，经过上述旋转，其父节点必为红；
                // 但如果 sibling 一开始就是 BLACK，父可红可黑。因此这里需要另外判断。
                if parent.Color == COLOR_RED {
                    parent.Color = COLOR_BLACK
                    break
                }

                // 如果父节点本来就是黑色，说明本次删除影响到祖父节点左右路径黑色节点数不一致，需要再往上看
                node = parent

                continue
            }

            /** 到这里时，sibling 左右必有一个红色节点，sibling 必为黑色 */

            // 由于 parent 左边要删除一个黑色节点，因此最终必须左旋保证 parent 左右路径黑色节点一致
            // 由于最终左旋后右子树黑色节点减少，因此需要先右旋增加右边黑色

            // sibling 右边为黑，则左边必为红。
            // 黑色的 node 是 parent 左，黑色的 sibling_right 是 sibling 右，左右相反，需要先反向旋转
            //               parent(?)
            //          /                  \
            //  node(BLACK)               sibling(BLACK)
            //     /    \           /                           \
            //   (?)    (?) sibling_left(RED)           sibling_right(BLACK|NULL_NODE)
            //                  /     \                         /    \
            // (?)(BLACK|NULL_NODE)  (?)(BLACK|NULL_NODE)    (?)     (?)
            if sibling.Right.Color == COLOR_BLACK {
                sibling.Left.Color = COLOR_BLACK
                sibling.Color = COLOR_RED

                // 红色顶点右旋，为下面的左旋做准备
                RotateRight(sibling)
                //                 parent(?)
                //          /                      \
                //  node(BLACK)              sibling_left(BLACK)
                //     /    \           /                           \
                //   (?)    (?)    (?)(BLACK|NULL_NODE)           sibling(RED)
                //                  /     \                       /    \
                //                (?)     (?)                   (?)    sibling_right(BLACK|NULL_NODE)

                sibling = parent.Right
                //                    parent(?)
                //          /                            \
                //  node(BLACK)                         sibling(BLACK)
                //     /    \                           /           \
                //   (?)    (?)  sibling_left(BLACK|NULL_NODE)       sibling_right(RED)
                //                  /     \                              /    \
                //                (?)     (?)                          (?)    (?)(BLACK|NULL_NODE)
            }

            // 父节点可红可黑
            sibling.Color = parent.Color
            parent.Color = COLOR_BLACK
            sibling.Right.Color = COLOR_BLACK
            //                 parent(BLACK)
            //          /                          \
            //  node(BLACK)                       sibling(?)
            //     /    \                      /               \
            //   (?)    (?)  sibling_left(BLACK|NULL_NODE)      sibling_right(BLACK)
            //                  /     \                           /    \
            //                (?)     (?)       (?)(BLACK|NULL_NODE)    (?)(BLACK|NULL_NODE)

            RotateLeft(parent)
            //                        sibling(?)
            //                 /                         \
            //     parent(BLACK)                           sibling_right(BLACK)
            //         /    \                                 /           \
            // node(BLACK) sibling_left(BLACK|NULL_NODE)     (?)        (?)(BLACK|NULL_NODE)
            //                 /     \
            //               (?)     (?)

            // 如果 sibling_left 为 NULL_NODE，这是删除 node 前的操作，后面需要把 node 删除掉。
            // 如果 sibling_left 为 BLACK，这是已经删除后的调整。
            break
        }

        /** 情况二：node 在 parent 右边 */
        if node == parent.Right {
            sibling = node.Parent.Left

            if sibling.Color == COLOR_RED {
                sibling.Color = COLOR_BLACK
                parent.Color = COLOR_RED
                RotateRight(parent)
                sibling = parent.Left
            }

            if sibling.Left.Color == COLOR_BLACK && sibling.Right.Color == COLOR_BLACK {
                sibling.Color = COLOR_RED

                if parent.Color == COLOR_RED {
                    parent.Color = COLOR_BLACK
                    break
                }

                node = parent

                continue
            }

            if sibling.Left.Color == COLOR_BLACK {
                sibling.Right.Color = COLOR_BLACK
                sibling.Color = COLOR_RED
                RotateLeft(sibling)
                sibling = parent.Left
            }

            sibling.Color = parent.Color
            parent.Color = COLOR_BLACK
            sibling.Left.Color = COLOR_BLACK
            RotateRight(parent)

            break
        }
    }

    tree.root.Color = COLOR_BLACK
}

// Successor 返回中序遍历下一个节点
func (tree *RedBlackTree) Successor(node *TreeNode) *TreeNode {
    if node == nil {
        return nil
    }

    if node.Right != nil {
        current := node.Right
        for current.Left != nil {
            current = current.Left
        }
        return current
    }

    for node.Parent != nil {
        if node.Parent.Left == node {
            return node.Parent
        }
        node = node.Parent
    }

    return nil
}

// Successor 返回中序遍历上一个节点
func (tree *RedBlackTree) Predecessor(node *TreeNode) *TreeNode {
    if node == nil {
        return nil
    }

    if node.Left != nil {
        current := node.Left
        for current.Right != nil {
            current = current.Right
        }
        return current
    }

    for node.Parent != nil {
        if node.Parent.Right == node {
            return node.Parent
        }
        node = node.Parent
    }

    return nil
}

// RotateLeft 左旋
func RotateLeft(root *TreeNode) *TreeNode {
    if root == nil {
        return nil
    }

    newRoot := root.Right
    newRoot.Parent = root.Parent
    // 红黑树旋转如果不在这里操作 newRoot 的 parent，
    // 某些情况就必须在调用该函数的地方做这个判断
    if root.Parent != nil {
        if root.Parent.Left == root {
            root.Parent.Left = newRoot
        } else {
            root.Parent.Right = newRoot
        }
    }
    root.Parent = newRoot

    root.Right = newRoot.Left
    newRoot.Left = root
    if root.Right != nil {
        root.Right.Parent = root
    }

    return newRoot
}

// RotateRight 右旋
func RotateRight(root *TreeNode) *TreeNode {
    if root == nil {
        return nil
    }

    newRoot := root.Left
    newRoot.Parent = root.Parent
    // 红黑树旋转如果不在这里操作 newRoot 的 parent，
    // 某些情况就必须在调用该函数的地方做这个判断
    if root.Parent != nil {
        if root.Parent.Left == root {
            root.Parent.Left = newRoot
        } else {
            root.Parent.Right = newRoot
        }
    }
    root.Parent = newRoot

    root.Left = newRoot.Right
    newRoot.Right = root
    if root.Left != nil {
        root.Left.Parent = root
    }

    return newRoot
}