手寫一棵紅黑樹
我記得面試的時候,經常問問別人hashmap
實現,說著說著就免不了講講紅黑樹,平常都是用現成的,考察別人紅黑樹也只是看下是否喜歡專研、有學習勁。
有一次有個同學告訴我他講不清楚但是可以寫一下,很慚愧,全忘了,一下子讓我寫一個,虛擬碼都夠嗆了,跑起來更不行。
我給自己想了個簡單的記法,父紅叔紅就變色,父紅叔黑靠旋轉,刪黑兩孩很麻煩,叔黑孩最很簡單。
–
紅黑樹
紅黑樹是AVL樹的進一步加強,正是二叉平衡查詢樹有問題才引出了紅黑樹,和典型資料結構一樣,在適當的場景使用紅黑樹可以很大程度的提高效能。
紅黑樹首先是一棵二叉查詢樹,節點的左孩子都比節點小,節點的右孩子都比節點大,與AVL平衡樹期望帶到的效果一樣,都想左右子樹的深度相差不要太大,儘量平衡,以便提供平均查詢效率。
先記住一下紅黑樹的以下幾個特性,不用急著回憶,後面程式碼寫著寫著自然就想起來了。
- 節點要麼是黑色要麼是紅色,根節點固定為黑色,葉子節點也固定為黑色(不關鍵特性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條邏輯線了。
- 叔叔節點為紅色,太簡單了,變個色即可
- 叔叔節點是紅色,那就要進行左右旋了,先理解單純的各種假設條件下的一次右旋,即可理解其他
情況一:叔叔節點是紅色
這個好理解的,接下來看下叔叔節點是黑色
情況二:叔叔節點是黑色,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