紅黑樹代碼實現及註釋

代碼參考了以下兩個鏈接：

• 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
}