如何写一颗60行的红黑树(in Haskell)

Posted 2021-03-10 candy99

如何用Haskell写一颗红黑树

同步更新于Candy?的新家

Candy?在上学期的数算课上学了红黑树，但是他一直没写过。

最近他入门了一下Haskell，得知用Haskell可以很方便实现各种树结构，于是就去学了一下如何用Haskell写红黑树，发现只要不到60行（包括空行和类型签名）！

下面是一个简单的小教程。

定义类型

和普通二叉树一样哒，只不过加上了一个颜色信息

data Tree a = Nil | Node Color (Tree a) a (Tree a) deriving (Show, Eq)
data Color = R | B deriving (Show, Eq)

辅助函数

• 将树根染黑：

makeBlack :: Tree a -> Tree a
makeBlack Nil = Nil
makeBlack (Node _ l x r) = Node B l x r

• 将树根染红：

makeRed :: Tree a -> Tree a
makeRed Nil = Nil
makeRed (Node _ l x r) = Node R l x r

插入操作

一般的红黑树插入不太方便用纯函数式来写，Okasaki在1999年提出了一种新的插入方法，将插入统一为：

• 首先默认插入红色节点，然后从下向上进行balance操作；
• balance操作会处理当前子树的children和grandchildren出现双红的情况，并且会将当前子树的根变红（balance操作并不会改变rank）

插入操作的框架很简单，需要注意的是最后要让整棵树的根变黑：

insert :: (Ord a) => a -> Tree a -> Tree a
insert x = makeBlack . ins 
  where ins Nil = Node R Nil x Nil
        ins t@(Node c l y r) | x < y     = balance $ Node c (ins l) y r
                             | x > y     = balance $ Node c l y (ins r)
                             | otherwise = t

balance操作要处理四种情况：

可以方便的用pattern matching来实现：

balance :: Tree a -> Tree a
balance (Node B (Node R (Node R a x b) y c) z d) = Node R (Node B a x b) y (Node B c z d)
balance (Node B (Node R a x (Node R b y c)) z d) = Node R (Node B a x b) y (Node B c z d)
balance (Node B a x (Node R (Node R b y c) z d)) = Node R (Node B a x b) y (Node B c z d)
balance (Node B a x (Node R b y (Node R c z d))) = Node R (Node B a x b) y (Node B c z d)
balance t@(Node c x l r) = t

删除操作

插入操作只要处理“双红”，删除操作还要处理“黑色节点数相等”，比较麻烦。

这里采用了Stefan Kahrs在2001年提出的方法，主要特点是：

• 不将带删除节点与后继交换
• 维持一个新的invariant：
  • 从黑根子树中删除节点，该子树高度会-1
  • 从红根子树中删除节点，该子树高度不变

我们有balanceL和balanceR两个操作，分别处理“左子树比右子树短1”和“右子树比左子树短1”的情况，将整棵树的高度变成较短那个的状态。

删除操作的框架如下：

delete :: Ord a => a -> Tree a -> Tree a
delete x = makeBlack . del
  where
    del Nil = Nil
    del t@(Node _ l y r) | x < y     = delL t
                         | x > y     = delR t
                         | otherwise = app l r
    delL (Node _ l@(Node B _ _ _) y r) = balanceL $ Node B (del l) y r
    delL (Node _ l y r)                = Node R (del l) y r
    delR (Node _ l y r@(Node B _ _ _)) = balanceR $ Node B l y (del r)
    delR (Node _ l y r)                = Node R l y (del r)

以待插入节点将插入左子树为例：

• 当前节点y的左子树为黑根时，会在删除后将y染黑并进行balanceL操作
• 当前节点y的左子树为红根时，会在删除后将y染红

容易发现，这样操作是可以维持新的invariant的（枚举当前节点颜色情况证明即可）

由于delete中在balanceL/R之前会染黑，balanceL/R只要处理根为黑的情况即可，有三种情况：

同样用pattern matching来实现：

balanceL :: Tree a -> Tree a 
balanceL (Node B (Node R a x b) y r) = Node R (Node B a x b) y r
balanceL (Node B l y (Node B a z b)) = balance $ Node B l y (Node R a z b)
balanceL (Node B l y (Node R (Node B a u b) z c)) = Node R (Node B l y a) u (balance $ Node B b z (makeRed c))

balanceR :: Tree a -> Tree a 
balanceR (Node B l y (Node R a x b)) = Node R l y (Node B a x b)
balanceR (Node B (Node B a z b) y r) = balance $ Node B (Node R a z b) y r
balanceR (Node B (Node R c z (Node B a u b)) y r) = Node R (balance $ Node B (makeRed c) z a) u (Node B b y r)

app会合并两个子树，有三种情况：

同样用pattern matching来实现：

app :: Tree a -> Tree a -> Tree a
app Nil t = t
app t Nil = t 
app (Node R a x b) (Node R c y d) = 
  case app b c of
    Node R b' z c' -> Node R (Node R a x b') z (Node R c' y d)
    s -> Node R a x (Node R s y d)
app (Node B a x b) (Node B c y d) =
  case app b c of
    Node r b' z c' -> Node R (Node B a x b') z (Node B c' y d)
    s -> balanceL $ Node B a x (Node B s y d)
app (Node R a x b) t = Node R a x (app b t)
app t (Node R a x b) = Node R (app t a) x b

完整代码

只要60行！

data Tree a = Nil | Node Color (Tree a) a (Tree a) deriving (Show, Eq)
data Color = R | B deriving (Show, Eq)

makeBlack :: Tree a -> Tree a
makeBlack Nil = Nil
makeBlack (Node _ l x r) = Node B l x r

makeRed :: Tree a -> Tree a
makeRed Nil = Nil
makeRed (Node _ l x r) = Node R l x r

insert :: (Ord a) => a -> Tree a -> Tree a
insert x = makeBlack . ins 
  where ins Nil = Node R Nil x Nil
        ins t@(Node c l y r) | x < y     = balance $ Node c (ins l) y r
                             | x > y     = balance $ Node c l y (ins r)
                             | otherwise = t

balance :: Tree a -> Tree a
balance (Node B (Node R (Node R a x b) y c) z d) = Node R (Node B a x b) y (Node B c z d)
balance (Node B (Node R a x (Node R b y c)) z d) = Node R (Node B a x b) y (Node B c z d)
balance (Node B a x (Node R (Node R b y c) z d)) = Node R (Node B a x b) y (Node B c z d)
balance (Node B a x (Node R b y (Node R c z d))) = Node R (Node B a x b) y (Node B c z d)
balance t@(Node c x l r) = t

delete :: Ord a => a -> Tree a -> Tree a
delete x = makeBlack . del
  where
    del Nil = Nil
    del t@(Node _ l y r) | x < y     = delL t
                         | x > y     = delR t
                         | otherwise = app l r
    delL (Node _ l@(Node B _ _ _) y r) = balanceL $ Node B (del l) y r
    delL (Node _ l y r)                = Node R (del l) y r
    delR (Node _ l y r@(Node B _ _ _)) = balanceR $ Node B l y (del r)
    delR (Node _ l y r)                = Node R l y (del r)

balanceL :: Tree a -> Tree a 
balanceL (Node B (Node R a x b) y r) = Node R (Node B a x b) y r
balanceL (Node B l y (Node B a z b)) = balance $ Node B l y (Node R a z b)
balanceL (Node B l y (Node R (Node B a u b) z c)) = Node R (Node B l y a) u (balance $ Node B b z (makeRed c))

balanceR :: Tree a -> Tree a 
balanceR (Node B l y (Node R a x b)) = Node R l y (Node B a x b)
balanceR (Node B (Node B a z b) y r) = balance $ Node B (Node R a z b) y r
balanceR (Node B (Node R c z (Node B a u b)) y r) = Node R (balance $ Node B (makeRed c) z a) u (Node B b y r)

app :: Tree a -> Tree a -> Tree a
app Nil t = t
app t Nil = t 
app (Node R a x b) (Node R c y d) = 
  case app b c of
    Node R b' z c' -> Node R (Node R a x b') z (Node R c' y d)
    s -> Node R a x (Node R s y d)
app (Node B a x b) (Node B c y d) =
  case app b c of
    Node r b' z c' -> Node R (Node B a x b') z (Node B c' y d)
    s -> balanceL $ Node B a x (Node B s y d)
app (Node R a x b) t = Node R a x (app b t)
app t (Node R a x b) = Node R (app t a) x b

其他API

一些其他常规操作的API：

tree2List :: Tree a -> [a]
tree2List Nil = []
tree2List (Node c l x r) = tree2List l ++ [x] ++ tree2List r

list2Tree :: Ord a => [a] -> Tree a
list2Tree = foldl (flip insert) Nil 

search :: (Ord a) => a -> Tree a -> Bool
search _ Nil = False
search x (Node _ l y r) 
  | x == y    = True
  | x < y     = search x l
  | otherwise = search x r

successor :: Ord a => a -> Tree a -> a
successor x Nil = x
successor x (Node _ l y r) 
  | x <  y = let t = successor x l in if x == t then y else t
  | x >= y = successor x r

PS：因为没有维护size信息所以没法求第k小QwQ，不过加上size信息应该也不难写。

参考资料

• 一篇讲的很好的博客
• 一份代码实现
• Kahrs, Stefan. (2001). Red-black trees with types. Journal of Functional Programming. 11. 10.1017/S0956796801004026.

另外，Matt Might提出了一种更加简洁、函数式的方法，详情参阅他的博客。