标签:haskell mamicode 教程 代码实现 代码 完整 Staff 枚举 入门
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年提出了一种新的插入方法,将插入统一为:
插入操作的框架很简单,需要注意的是最后要让整棵树的根变黑:
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年提出的方法,主要特点是:
我们有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)
以待插入节点将插入左子树为例:
容易发现,这样操作是可以维持新的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:
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提出了一种更加简洁、函数式的方法,详情参阅他的博客。
标签:haskell mamicode 教程 代码实现 代码 完整 Staff 枚举 入门
原文地址:https://www.cnblogs.com/candy99/p/rbt_in_haskell.html