手寫一棵紅黑樹

筆者博客地址：https://charpty.com

我記得面試的時候，經常問問別人hashmap實現，說着說着就免不了講講紅黑樹，平常都是用現成的，考察別人紅黑樹也只是看下是否喜歡專研、有學習勁。

有一次有個同學告訴我他講不清楚但是可以寫一下，很慚愧，全忘了，一下子讓我寫一個，僞代碼都夠嗆了，跑起來更不行。

我給自己想了個簡單的記法，父紅叔紅就變色，父紅叔黑靠旋轉，刪黑兩孩很麻煩，叔黑孩最很簡單。

–

紅黑樹

紅黑樹是AVL樹的進一步加強，正是二叉平衡查找樹有問題才引出了紅黑樹，和典型數據結構一樣，在適當的場景使用紅黑樹可以很大程度的提高性能。

紅黑樹首先是一棵二叉查找樹，節點的左孩子都比節點小，節點的右孩子都比節點大，與AVL平衡樹期望帶到的效果一樣，都想左右子樹的深度相差不要太大，儘量平衡，以便提供平均查找效率。

先記住一下紅黑樹的以下幾個特性，不用急着回憶，後面代碼寫着寫着自然就想起來了。

1. 節點要麼是黑色要麼是紅色，根節點固定爲黑色，葉子節點也固定爲黑色（不關鍵特性3合一）
2. 子節點和父節點不能同時爲紅色，子父不連紅。
3. 從一個節點到其通向到所有葉子節點路徑中，所包含的黑色節點數目相同。保證樹平衡的關鍵。

前面兩點都很好理解，第2點是用來修改樹時判斷樹是否還是紅黑樹的主要條件。
第3點不直觀，但是可以這樣想，插入或刪除一個節點，影響的只是它周邊那幾個節點（之外的節點本來就是「平衡」的），所以這句話可以翻譯成說，要在修改節點後，要把上、左、右這幾個位置上的黑色節點數量控制住，所以此時只要把周邊幾個節點挪一挪，就又恢復平衡了。

所以在紅黑樹實現中，一般不直接判斷第3點（一層層遍歷下去效率太低），而僅僅是把周圍幾個節點通過變色和旋轉來達到平衡。

對於紅黑樹的理論講解，網上非常多，但是我想實在點，直接一起寫吧，寫本文之前，我也是照着算法僞代碼直接開寫，很多忘了的都想起來了。

插入節點Z

和業務代碼一樣，紅黑樹也無非是增、刪、改、查，其它三個都包含着查，增和刪對樹結構變化最大，我們就看這兩個即可理解紅黑樹了，先來看插入節點的僞代碼（網上找了個，不太對我改了下）。

// 在插入節點（二叉查找樹的插入）完成後，如果破壞了紅黑樹特性，則對紅黑樹進行修復
// T表示當前紅黑樹，z表示當前插入的節點，->p表示父節點，->right表示右孩子，類推
RB-INSERT-FIXUP(T, z)
// 爲了不與「特性3」衝突，所以插入的z是紅色，這樣黑色節點的數目肯定是不會變化的
// 如果z的父節點爲紅那就與「特性2：子父節點不同時爲紅」衝突，此時要分幾種情況調整
while z->p->color = RED; do
	 // 如果z的父節點爲爺爺節點的左孩子
	if z->p = z->p->p->left then
 			y ← z->p->p->right
 			// 叔叔節點爲紅色或黑色，分爲兩種情況處理
 			if y->color = RED then
				// 如果叔叔是紅色，爺爺節點是黑色，這種情況比較簡單，此時無論父節點是爺爺節點的左還是右節點
				// 都是將父節點設置爲黑色，叔叔節點設置爲黑色，祖父節點設置爲紅色
				// 這樣一來，子父爲紅-紅的情況自然是不存在了，父節點和叔叔節點由紅-紅變成了黑-黑
				// 經過這兩個節點的到根節點路徑黑色節點數沒變，都是增加了一個黑色節點
				// 經過爺爺節點到根路徑的黑色節點數量則無變化，爺爺節點變成了紅色，但是它的兩個孩子不論選哪條路都加1了
				z->p->color ← BLACK
				y->color ← BLACK
				z->p->p->color ← RED
				// 爺爺節點設置爲紅色之後，繼續向上判斷它和其父節點是否衝突
				z ← z->p->p
 			else 
 				// 如果叔叔節點是黑色就需要旋轉樹了，如果x爲父節點的左孩子，先要額外進行一次進行左旋
 				if z = z->p->right then	
 					z ← z->p
 					LEFT-ROTATE(T, z)
 				// 先假設x爲父節點的左節點，這樣比較簡單，弄清楚了加一層左旋一樣的道理
 				z->p->color ← BLACK
 				z->p->p->color ← RED
 				// 上面兩行代碼已解決了"子父節點不能同爲紅色"的問題，這樣經過爺爺節點走左邊的話黑色節點計數還是不變的
                // 但是原本通過爺爺節點走右邊的話有兩個黑節點的，現在只有一個了，此時只有一個了
                // 關鍵來了，在節點爲紅-黑-紅-黑（頂上爲紅）的情況下，右旋使得旋轉節點的右孩子路徑上黑色節點數加1
 				RIGHT-ROTATE(T, z->p->p)
 	// 如果z的父親爲爺爺節點的右孩子，叔叔節點爲紅色的邏輯是一樣的，只是叔叔爲黑時邏輯「相反」
	else (same as then clause with "right" and "left" exchanged)
 T->root->color ← BLACK

爲了寫的更清楚，特地將Java的TreeMap又看了一遍，其中的fixAfterInsertion()函數正是這個邏輯。

到底幹了啥呢，其實就當兩種情況來理解的話，就沒那麼繞了。只是外面套了一層父節點是爺爺節點的左還是右節點，導致2*2變成4條邏輯線了。

1. 叔叔節點爲紅色，太簡單了，變個色即可
2. 叔叔節點是紅色，那就要進行左右旋了，先理解單純的各種假設條件下的一次右旋，即可理解其他

情況一：叔叔節點是紅色

這個好理解的，接下來看下叔叔節點是黑色

情況二：叔叔節點是黑色，Z的父節點爲爺爺節點的左孩子，Z也爲父節點左孩子

原來的邏輯是先塗色，再右旋，但是不能很好的體現左旋的作用，不管是左旋還是右旋，邏輯都是將紅色節點向根節點靠攏，最後將紅色節點塗黑。

也就是以下流程

情況三：叔叔節點是黑色，Z的父節點爲爺爺節點的左孩子，Z也爲父節點右孩子

此時就比較麻煩了，處理的思路是將情況三轉換爲情況二，這需要額外的一次左旋。

可以看到，情況三是先把問題轉化爲情況二，再利用已知的處理方式調整

還有另外一個邏輯和情況二、三相反，就不重複敘述了。

刪除節點X

和插入的邏輯類似，插入時是先按照二叉查找樹的方式先插入再調整，刪除時也是先按照二叉查找樹的方式先刪除，然後再調整。

要提醒的是，二叉查找樹的刪除，不論刪除哪個節點，最終都是刪除「最邊上」的節點，要麼是葉子節點，要麼是有一個孩子的節點，度最大爲1。因爲即使刪除中間的某個節點，也得選它左子樹中最大的節點補上去（選左右都一樣），那左子樹最大的節點肯定是在左子樹右邊「最邊上」了。

和二叉查找樹稍有不同的是，紅黑樹是帶顏色的，爲了保證「上邊」的樹結構滿足紅黑樹特性，所以補上節點時，僅僅是把節點的值拷貝過去，顏色不拷貝。

所以接下來我們討論的都是刪除這個「最邊上」節點的種種情況，稱之爲X節點。

刪除操作的僞代碼

// 在刪除節點操作完成後對紅黑樹進行修復
RB-DELETE-FIXUP(T, x)
// 刪root沒啥好處理的，刪紅色節點也無需理會（後續有講解爲何）
while x ≠ root[T] and color[x] = BLACK  do
	// 在寫僞代碼以及操作解釋時都僅說明x爲父節點左孩子的情況，右孩子情況是對稱的
	if x = left[p[x]] then  
		// 關注的是x的兄弟節點和其孩子節點的情況 
		w ← right[p[x]]  
		// 兄弟節點是紅色，則將其轉換爲"兄弟節點是黑色"的情況
		if color[w] = RED  then 
			color[w] ← BLACK                         
			color[p[x]] ← RED                         
			LEFT-ROTATE(T, p[x])                      
			w ← right[p[x]]                           
		if color[left[w]] = BLACK and color[right[w]] = BLACK  then 
			// 兄弟節點及其孩子節點均爲黑色的情況下，則將其轉換爲"兄弟節點爲紅色"
			color[w] ← RED                            
			x ← p[x]                                  
		else 
			if color[right[w]] = BLACK then
				// 直接轉換爲"兄弟節點右孩子爲紅色"情況
				color[left[w]] ← BLACK            
				color[w] ← RED                   
				RIGHT-ROTATE(T, w)                
				w ← right[p[x]]
			// 兄弟節點右孩子爲紅色的情況可以一步到位達到平衡                   
			color[w] ← color[p[x]]                   
			color[p[x]] ← BLACK                      
			color[right[w]] ← BLACK                  
			LEFT-ROTATE(T, p[x])                   
			x ← root[T]                             
	else (same as then clause with "right" and "left" exchanged)  
color[x] ← BLACK

這裏容易混淆的是，比如以A爲中心左旋時，A成爲A的右孩子的左孩子，A的右孩子的左孩子B成爲A的右孩子，注意B在成爲A的右孩子時，是將B以及B下面整棵子樹娜過來了。

比較簡單的幾種情況

刪除的節點X是紅色
如果刪除的節點X是紅色，那麼首先說明原來上下都是黑色的，刪了X節點一不違背「子父節點不同時爲紅」的特性，二不違背「各節點到葉節點路徑上黑色節點數目相同」的特性，所以無需處理。

接替X的節點W是紅色
X被刪了，它自己是一個黑色，它的子節點有且僅有一個，顏色是紅色。
接替它的節點W是紅色，那麼直接用W接替X的位置，再把W塗黑即可。

X爲根節點的情況
如果X爲黑色，W也是黑色，那就比較麻煩了，分很多情況，其中最特殊的就是X是根節點，此時刪除X之後啥也不用做，刪除根節點唯一要考慮的僅僅是「紅-紅」衝突而已。

另外的情況比較複雜，每種情況的處理方式都不同，我們僅舉X是其父節點的左孩子的情況，和插入一下，爲右孩子時，操作是對稱的。

值得注意的是如果X是黑色且沒有任何子節點，那麼也是通過旋轉等複雜操作來重新平衡的，這時我們就假設替代的節點是個黑色節點（虛擬的）就行，主要看的是X的兄弟以及X的侄子的顏色情況。

情況一：X、W是黑色，X的兄弟節點Y是黑色，Y的右孩子是紅色

前面3個條件的處理的都是最簡單的情況，我們當然希望要刪除的都是紅色，這樣啥也不用幹了，但是接下來4種情況都是比較繞的。

雖然複雜，但是記住一個原則，後面的情況二、情況四、情況五所做的動作，都是想最終轉化爲情況一或上述3種簡單情況而已。也就是說，情況一和上面的3種情況是與紅黑樹平衡最接近的場景，只需一步操作即可恢復平衡了，而其他情況則需要先轉換爲這些情況。

做法也還是塗色加旋轉，先把兄弟節點Y染成當X的父節點的顏色，再把X節點父節點染成黑色，Y節點右孩子子染成黑色，最後再以X節點的父節點爲中心進行左旋。

爲了和後面的情況統一風格，我們認定情況一的處理辦法爲：

情況一  ->  最終平衡

情況二：X、W是黑色，X的兄弟點Y是黑色，Y的右孩子爲黑色，Y的左孩子爲紅色

此時我們要做的事將該場景轉換爲情況一，然後我們再使用情況一的解決辦法即可。

做法是將兄弟節點Y塗紅，Y節點左孩子塗黑，之後再以兄弟節點Y爲中心右旋。

這種情況的處理辦法爲

情況二  ->  情況一  ->  最終平衡

情況三：X、W是黑色，X的兄弟Y爲紅色

此時X的父節點以及Y的孩子均爲黑色，處理原則是將X的兄弟節點變爲黑色（當然是在不能破壞目前的紅黑樹已有性質前提下）。

具體處理辦法是以X的父節點爲中心進行左旋。左旋之後X的新兄弟節點必然爲黑色，此時又回到了兄弟節點爲黑色的幾種情況上。

這種情況的處理辦法爲

情況三  -> （情況一、情況二）

情況四：X、W是黑色，X的父親、Y及其孩子均爲黑色

這種情況下，左邊X的路徑上因爲刪除少了一個黑色節點，此時我們將Y節點塗紅，這樣經過Y和經過W（替代X後）的黑色節點數達到一致了。

但問題是經過原X的父節點的路徑的黑色節點數少1了，但此時整個結構又回到了情況四（右邊路徑上黑色數目不同了但不影響），所以我們又可以按照情況四繼續往下走。

這種情況的處理辦法爲

情況四  -> 情況三

紅黑樹實現

想了想還是用Python寫吧，人生苦短。

完整的代碼：請下載

定義一個class表示紅黑樹吧，只要存一個root節點就夠了。

class RBTree:
    def __init__(self):
        self.root = None

    def insert(self, key, value):
        if self.root is None:
            self.root = Node(key, value, BLACK, None)
            return
        parent = None
        t = self._get_node(key)
        if t is not None:
            t.value = value
            return
        node = Node(key, value, RED, parent)
        if parent.key < node.key:
            parent.right = node
        else:
            parent.left = node

        fix_insert(node, self)

    def delete(self, key):
        x = self._get_node(key)
        if x is None:
            return
        if left_child(x) is not None and right_child(x) is not None:
            real_delete = get_successor(x)
            x.key = real_delete.key
            x.value = real_delete.value
            x = real_delete
        # 到此，x最多也就一個孩子了
        successor = get_one_child(x)
        if successor is None:
            self._delete_leaf(x)
            return

        if get_parent(x) is None:
            self.root = None
        elif x is left_child(get_parent(x)):
            get_parent(x).left = successor
        else:
            get_parent(x).right = successor

        if get_color(x) is BLACK:
            fix_delete(successor, self)

插入修復

def fix_insert(node, tree):
    z = node
    while z is not None and z is not tree.root and get_color(get_parent(z)) is RED:
        if get_parent(z) is left_child(get_grandparent(z)):
            uncle = right_child(get_grandparent(z))
            if get_color(uncle) is RED:
                # 叔叔是紅色，此時將父親和叔叔設置爲黑色，爺爺設置爲紅色即可
                # 不管父節點是左孩子還是右孩子都一樣
                set_color(get_parent(z), BLACK)
                set_color(uncle, BLACK)
                set_color(get_grandparent(z), RED)
                # 爺爺節點塗紅之後，繼續向上同樣方式判斷
                z = get_parent(z)
            else:
                # 如果z爲父節點的右孩子，先要把它變成左孩子形式（實質上父子對掉了）
                if z is right_child(get_parent(z)):
                    z = get_parent(z)
                    rotate_left(z, tree)
                # 此時z爲父節點的左孩子，將父節點塗黑，爺爺節點塗紅，在以爺爺節點爲中心右旋
                set_color(get_parent(z), BLACK)
                set_color(get_grandparent(z), RED)
                rotate_right(get_grandparent(z), tree)
        else:
            uncle = left_child(get_grandparent(z))
            if get_color(uncle) is RED:
                set_color(get_parent(z), BLACK)
                set_color(uncle, BLACK)
                set_color(get_grandparent(z), RED)
                z = get_parent(z)
            else:
                if z is left_child(get_parent(z)):
                    z = get_parent(z)
                    rotate_right(z, tree)
                set_color(get_parent(z), BLACK)
                set_color(get_grandparent(z), RED)
                rotate_left(get_grandparent(z), tree)

刪除修復

def fix_delete(node, tree):
    x = node
    while x is not tree.root and get_color(x) is BLACK:
        if x is left_child(get_parent(x)):
            brother = right_child(get_parent(x))
            if get_color(brother) is RED:
                # 兄弟節點是紅色，則將其轉換爲"兄弟節點是黑色"的情況
                set_color(brother, BLACK);
                set_color(get_parent(x), RED);
                rotate_left(get_parent(x));
                brother = right_child(get_parent(x));
            if get_color(left_child(brother)) is BLACK and get_color(right_child(brother)) is BLACK:
                # 兄弟節點及其孩子節點均爲黑色的情況下，則將其轉換爲"兄弟節點爲紅色"
                set_color(brother, RED)
                x = get_parent(x)
            else:
                if get_color(right_child(brother)) is BLACK:
                    # 直接轉換爲"兄弟節點右孩子爲紅色"情況
                    set_color(left_child(brother), BLACK)
                    set_color(brother, RED)
                    rotate_right(brother, tree)
                    brother = right_child(get_parent(x))
                # 這裏僅一步即可達到平衡
                set_color(brother, get_color(get_parent(x)))
                set_color(get_parent(x), BLACK)
                set_color(right_child(brother), BLACK)
                rotate_left(get_parent(x), tree)
                x = tree.root
        else:
            brother = left_child(get_parent(x))
            if get_color(brother) is RED:
                set_color(brother, BLACK);
                set_color(get_parent(x), RED);
                rotate_right(get_parent(x));
                brother = left_child(get_parent(x));
            if get_color(right_child(brother)) is BLACK and get_color(left_child(brother)) is BLACK:
                set_color(brother, RED)
                x = get_parent(x)
            else:
                if get_color(left_child(brother)) is BLACK:
                    set_color(right_child(brother), BLACK)
                    set_color(brother, RED)
                    rotate_left(brother, tree)
                    brother = left_child(get_parent(x))
                set_color(brother, get_color(get_parent(x)))
                set_color(get_parent(x), BLACK)
                set_color(left_child(brother), BLACK)
                rotate_right(get_parent(x), tree)
                x = tree.root