zoukankan      html  css  js  c++  java
  • 如何写一颗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提出了一种更加简洁、函数式的方法,详情参阅他的博客

  • 相关阅读:
    UVA 11174 Stand in a Line,UVA 1436 Counting heaps —— (组合数的好题)
    UVA 1393 Highways,UVA 12075 Counting Triangles —— (组合数,dp)
    【Same Tree】cpp
    【Recover Binary Search Tree】cpp
    【Binary Tree Zigzag Level Order Traversal】cpp
    【Binary Tree Level Order Traversal II 】cpp
    【Binary Tree Level Order Traversal】cpp
    【Binary Tree Post order Traversal】cpp
    【Binary Tree Inorder Traversal】cpp
    【Binary Tree Preorder Traversal】cpp
  • 原文地址:https://www.cnblogs.com/candy99/p/rbt_in_haskell.html
Copyright © 2011-2022 走看看