RedBlackTree.hs (3978B)
1 {-# OPTIONS_GHC -Wall #-} 2 3 module RedBlackTree where 4 5 data Color = Red | Black deriving (Eq, Show) 6 7 data Tree a = Leaf | Node Color (Tree a) a (Tree a) deriving Show 8 9 empty :: Tree a 10 empty = Leaf 11 12 isEmpty :: Tree a -> Bool 13 isEmpty Leaf = True 14 isEmpty _ = False 15 16 member :: Ord a => a -> Tree a -> Bool 17 member _ Leaf = False 18 member e0 (Node _ l e1 r) = case compare e0 e1 of 19 LT -> member e0 l 20 GT -> member e0 r 21 EQ -> True 22 23 singleton :: a -> Tree a 24 singleton e = Node Red Leaf e Leaf 25 26 blacken :: Tree a -> Tree a 27 blacken (Node _ l e r) = Node Black l e r 28 blacken Leaf = Leaf 29 30 insert :: Ord a => a -> Tree a -> Tree a 31 insert e0 = blacken . ins where 32 ins Leaf = singleton e0 33 ins t@(Node c l e1 r) = case compare e0 e1 of 34 LT -> balance c (ins l) e1 r 35 GT -> balance c l e1 (ins r) 36 EQ -> t 37 38 balance :: Color -> Tree a -> a -> Tree a -> Tree a 39 balance Black (Node Red (Node Red a x b) y c) z d = 40 Node Red (Node Black a x b) y (Node Black c z d) 41 42 balance Black (Node Red a x (Node Red b y c)) z d = 43 Node Red (Node Black a x b) y (Node Black c z d) 44 45 balance Black a x (Node Red (Node Red b y c) z d) = 46 Node Red (Node Black a x b) y (Node Black c z d) 47 48 balance Black a x (Node Red b y (Node Red c z d)) = 49 Node Red (Node Black a x b) y (Node Black c z d) 50 51 balance c l e r = Node c l e r 52 53 -- exercise 3.8 (prove max depth of node is 2*floor(log (n + 1))) 54 -- - empty nodes black 55 -- - invariant 1: no red node has a red child 56 -- - invariant 2: every path from root to leaf contains same number of black nodes 57 58 -- exercise 3.9 (fromOrdList) 59 fromOrdList :: Ord a => [a] -> Tree a 60 fromOrdList = foldr insert empty 61 62 -- exercise 3.10a (improve balance) 63 lbalance :: Color -> Tree a -> a -> Tree a -> Tree a 64 lbalance Black (Node Red (Node Red a x b) y c) z d = 65 Node Red (Node Black a x b) y (Node Black c z d) 66 67 lbalance Black (Node Red a x (Node Red b y c)) z d = 68 Node Red (Node Black a x b) y (Node Black c z d) 69 70 lbalance c l e r = Node c l e r 71 72 rbalance :: Color -> Tree a -> a -> Tree a -> Tree a 73 rbalance Black a x (Node Red (Node Red b y c) z d) = 74 Node Red (Node Black a x b) y (Node Black c z d) 75 76 rbalance Black a x (Node Red b y (Node Red c z d)) = 77 Node Red (Node Black a x b) y (Node Black c z d) 78 79 rbalance c l e r = Node c l e r 80 81 altInsert :: Ord a => a -> Tree a -> Tree a 82 altInsert e0 = blacken . ins where 83 ins Leaf = singleton e0 84 ins t@(Node c l e1 r) = case compare e0 e1 of 85 LT -> lbalance c (ins l) e1 r 86 GT -> rbalance c l e1 (ins r) 87 EQ -> t 88 89 altFromList :: Ord a => [a] -> Tree a 90 altFromList = foldr altInsert empty 91 92 -- exercise 3.10b (improve balance) 93 lrbalance :: Color -> Tree a -> a -> Tree a -> Tree a 94 lrbalance Black (Node Red a x (Node Red b y c)) z d = 95 Node Red (Node Black a x b) y (Node Black c z d) 96 97 lrbalance c a x b = Node c a x b 98 99 llbalance :: Color -> Tree a -> a -> Tree a -> Tree a 100 llbalance Black (Node Red (Node Red a x b) y c) z d = 101 Node Red (Node Black a x b) y (Node Black c z d) 102 103 llbalance c a x b = Node c a x b 104 105 rlbalance :: Color -> Tree a -> a -> Tree a -> Tree a 106 rlbalance Black a x (Node Red (Node Red b y c) z d) = 107 Node Red (Node Black a x b) y (Node Black c z d) 108 109 rlbalance c a x b = Node c a x b 110 111 rrbalance :: Color -> Tree a -> a -> Tree a -> Tree a 112 rrbalance Black a x (Node Red b y (Node Red c z d)) = 113 Node Red (Node Black a x b) y (Node Black c z d) 114 115 rrbalance c a x b = Node c a x b 116 117 altAltInsert :: Ord a => a -> Tree a -> Tree a 118 altAltInsert e0 s = blacken $ ins s where 119 ins Leaf = singleton e0 120 ins (Node c l e1 r) = case compare e0 e1 of 121 LT -> case l of 122 Leaf -> Node c (ins l) e1 r 123 Node _ _ e2 _ 124 | e0 <= e2 -> llbalance c (ins l) e1 r 125 | otherwise -> lrbalance c (ins l) e1 r 126 127 GT -> case r of 128 Leaf -> Node c l e1 (ins r) 129 Node _ _ e2 _ 130 | e0 <= e2 -> rlbalance c l e1 (ins r) 131 | otherwise -> rrbalance c l e1 (ins r) 132 133 EQ -> s 134 135 altAltFromList :: Ord a => [a] -> Tree a 136 altAltFromList = foldr altAltInsert empty 137