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

如何用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操作要处理四种情况:

rbt1

可以方便的用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只要处理根为黑的情况即可,有三种情况:

rbt2同样用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会合并两个子树,有三种情况:

rbt3

同样用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信息应该也不难写。

参考资料

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

posted @ 2020-02-11 17:55  Candy?  阅读(1213)  评论(0编辑  收藏  举报