Implemention of Left Leaning Red Black Tree, see http://www.cs.princeton.edu/~rs/talks/LLRB/RedBlack.pdf

LLRBTree.hs

{-# LANGUAGE BangPatterns #-}
module LLRBTree where
 
import Data.List (foldl', nub, (\\))
import qualified Data.List as L
import Data.Maybe
import System.Cmd
import Prelude hiding (minimum)
import Debug.Trace
import Test.QuickCheck
 
data LLRBTree a = Leaf
                | Node Color (LLRBTree a) a (LLRBTree a)
                deriving (Eq, Show)

isRed (Node Red _ _ _) = True
isRed _ = False

left Leaf = Leaf
left (Node _ l _ _) = l

right Leaf = Leaf
right (Node _ _ _ r) = r
 
data Color = Black
           | Red
           deriving (Eq, Show)
 
rotateLeft (Node c l v r@(Node Red rl rv rr)) 
    = Node c (Node Red l v rl) rv rr
rotateLeft _ = error "rotateLeft"
 
rotateRight (Node c l@(Node Red ll lv lr) v r)
    = Node c ll lv (Node Red lr v r)
rotateRight _ = error "rotateRight"
 
colorFlipUp (Node Black l@(Node Red _ _ _) v r@(Node Red _ _ _))
    = Node Red (turnB l) v (turnB r)
colorFlipUp t = t
 
insert t v = turnB (insert' t v)
 
insert' Leaf v = Node Red Leaf v Leaf
insert' (Node c l v r) x = case compare x v of
    LT -> fixUp (Node c (insert' l x) v r)
    GT -> fixUp (Node c l v (insert' r x))
    EQ -> Node c l x r
 
fixUp = colorFlipUp . fixL2L . fixRL

{- fix right leaning -}
fixRL t@(Node _ _ _ (Node Red _ _ _)) = rotateLeft t
fixRL t = t

{- fix left left leaning -}
fixL2L t@(Node _ (Node Red (Node Red _ _ _) _ _) _ _) = rotateRight t
fixL2L t = t
 
turnB (Node c l v r) = Node Black l v r
turnB Leaf = Leaf

turnR (Node c l v r) = Node Red l v r
turnR Leaf = Leaf

delMax Leaf = Leaf
delMax t = turnB (delMax' (turnR t))

delMax' (Node Red Leaf _ Leaf) = Leaf
-- right child's left child is red
delMax' (Node c l v r@(Node Black rl@(Node Red _ _ _) rv rr))
    = fixUp (Node c l v (delMax' r))
-- left child is red
delMax' t@(Node Black l@(Node Red _ _ _) v r) = fixUp (Node c' l' v' (delMax' r'))
  where
    (Node c' l' v' r') = rotateRight t
-- node is red and left child's left child is black or leaf
delMax' (Node Red l@(Node Black ll@(Node Black _ _ _) _ _) v r)
    = fixUp (Node Black (turnR l) v (delMax' (turnR r)))
delMax' (Node Red l@(Node Black Leaf _ _) v r)
    = fixUp (Node Black (turnR l) v (delMax' (turnR r)))
-- node is red and left child's left child is red
delMax' (Node Red l@(Node Black ll@(Node Red _ _ _) _ _) v r)
    = fixUp (Node c' l' v' (Node rc' rl' rv' (delMax' rr')))
  where
    (Node c' l' v' (Node rc' rl' rv' rr')) = Node Red (turnB l'') v'' (turnB r'')
    (Node c'' l'' v'' r'') = rotateRight (Node Black (turnR l) v (turnR r))
delMax' _ = error "delMax'"

delMin Leaf = Leaf
delMin t = turnB (delMin' (turnR t))
delMin' (Node Red Leaf _ Leaf) = Leaf
delMin' (Node c l@(Node Red _ _ _) v r) = fixUp (Node c (delMin' l) v r)
delMin' (Node c l@(Node _ (Node Red _ _ _) _ _) v r)
    = fixUp (Node c (delMin' l) v r)
delMin' (Node Red l v r@(Node _ (Node Black _ _ _) _ _))
    = fixUp (Node Black (delMin' (turnR l)) v (turnR r))
delMin' (Node Red l v r@(Node _ Leaf _ _))
    = fixUp (Node Black (delMin' (turnR l)) v (turnR r))
delMin' (Node Red l v r@(Node _ (Node Red _ _ _) _ _))
    = fixUp (Node c' (delMin l') v' r')
  where
    (Node c' l' v' r') = Node Red (turnB l'') v'' (turnB r'')
    (Node c'' l'' v'' r'') = rotateLeft (Node Black (turnR l) v (rotateRight (turnR r)))
delMin' _ = error "delMin'"

moveRedRight t = (rotateR . colorFlip) t
  where
    rotateR t@(Node _ (Node _ (Node Red _ _ _) _ _) _ _) = (colorFlip . rotateRight) t
    rotateR t = t

condMoveRedRight t
    | isRbRLb = moveRedRight t
    | otherwise = t
  where
    isRbRLb = (not . isRed . right) t && (not . isRed . left . right) t

deleteMax Leaf = Leaf
deleteMax t = turnB (deleteMax' (turnR t))

deleteMax' (Node Red Leaf _ Leaf) = Leaf
deleteMax' t
    | isLr = (fixUp . deleteMaxR . condMoveRedRight . rotateRight) t
    | otherwise = (fixUp . deleteMaxR . condMoveRedRight) t
  where
    isLr = (isRed . left) t
    deleteMaxR (Node c l v r) = Node c l v (deleteMax' r)

colorFlip (Node Black (Node Red ll lv lr) v (Node Red rl rv rr))
    = Node Red (Node Black ll lv lr) v (Node Black rl rv rr)
colorFlip (Node Red (Node Black ll lv lr) v (Node Black rl rv rr))
    = Node Black (Node Red ll lv lr) v (Node Red rl rv rr)
colorFlip t = error ("colorFlip (" ++ show t ++ ")")

deleteMin Leaf = Leaf
deleteMin t = turnB $ deleteMin' $ turnR $ t

deleteMin' (Node Red Leaf _ Leaf) = Leaf
deleteMin' t = (fixUp . deleteMinL . condMoveRedLeft) t
  where
    deleteMinL (Node c l v r) = Node c (deleteMin' l) v r

condMoveRedLeft t
    | isLbLLb = moveRedLeft t
    | otherwise = t
  where
    isLbLLb = (not . isRed .left) t && (not . isRed . left . left) t
    
moveRedLeft t = (rotateL . colorFlip) t
  where
    rotateL t@(Node c l v r@(Node _ (Node Red _ _ _) _ _)) = (colorFlip . rotateLeft) (Node c l v (rotateRight r))
    rotateL t = t

delete t k = turnB (delete' (turnR t) k)

delete' Leaf _ = Leaf
delete' t@(Node c l v r) k
    | k < v  = deleteLT t k
    | k == v = deleteEQ t k
    | k > v  = deleteGT t k

deleteLT t k
    | isLbLLb   = (fixUp . deleteL k . moveRedLeft) t
    | otherwise = (fixUp . deleteL k) t
  where
    isLbLLb = (not . isRed .left) t && (not . isRed . left . left) t
    deleteL k (Node c l v r) = Node c (delete' l k) v r

deleteEQ (Node Red Leaf _ Leaf) _ = Leaf
deleteEQ t k
    | isLr = (fixUp . deleteR k . rotateRight) t
    | isRbRLb && isLbLLr = (fixUp . deleteR k . moveRedRight) t
    | isRbRLb = (fixUp . deleteM . moveRedRight) t
    | otherwise = (fixUp . deleteM) t
  where
    isLr    = (isRed . left) t
    isRbRLb = (not . isRed . right) t && (not . isRed . left . right) t
    isLbLLr = (not . isRed . left) t && (isRed . left . left) t
    deleteM t@(Node _ Leaf _ Leaf) = Leaf
    deleteM t@(Node c l v r) = Node c l (minimum r) (deleteMin' r)
    deleteR k (Node c l v r) = Node c l v (delete' r k)

deleteGT t k
    | isLr = (fixUp . deleteR k . rotateRight) t
    | isRbRLb = (fixUp . deleteR k . moveRedRight) t
    | otherwise = (fixUp . deleteR k) t
  where
    isLr    = (isRed . left) t
    isRbRLb = (not . isRed . right) t && (not . isRed . left . right) t
    deleteR k (Node c l v r) = Node c l v (delete' r k)

leanRight t@(Node c (Node Red _ _ _) _ _) = rotateRight t
leanRight t = t

member v1 Leaf = False
member v1 t@(Node _ l v r)
    | v1 == v = True
    | v1 < v = member v1 l
    | v1 > v = member v1 r

minimum (Node _ Leaf v _) = v
minimum (Node _ l _ _) = minimum l

{- tool functions -}
 
empty = Leaf 
 
fromList = foldl' insert empty

toList t = inorder t []
  where
    inorder Leaf xs = xs
    inorder (Node _ l x r) xs = inorder l (x : inorder r xs)
 
 
{------------- do some test --------------}
 
dotTree t = "digraph {\nnode [style = filled];\n" ++ nodes ++ "\n}\n"
  where
    (_,_,nodes) = dotTree' t
 
dotTree' Leaf = (Black, Nothing, "")
dotTree' (Node c l v r) = (c, Just n, n_attr ++ nl ++ nr)
  where
    n = show v
    n_attr = if c == Red then n ++ "[fillcolor = red];\n" else "\n"
    (lc, ln, lt) = dotTree' l
    (rc, rn, rt) = dotTree' r
    nl = fromMaybe "" (ln >>= (\x -> Just (n ++ " -> " ++ x ++ ";\n"))) ++ lt
    nr = fromMaybe "" (rn >>= (\x -> Just (n ++ " -> " ++ x ++ ";\n"))) ++ rt
    
 
 
a = Node Black Leaf "A-E" Leaf
a_e = Node Black Leaf "A-E" Leaf
g_p = Node Black Leaf "G-P" Leaf
r_z = Node Black Leaf "R-Z" Leaf
q = Node Red g_p "Q" r_z
f = Node Black a_e "F" q
 
m = Node Black (Node Red Leaf "E" Leaf) "M" (Node Red Leaf "Q" Leaf)
 
test = f == (rotateRight . rotateLeft) f
 
demo = foldl' insert Leaf [100,98..0]
 
main = do
    writeFile fp1 dotScript1
    writeFile fp2 dotScript2
    rawSystem "dot" ["-Tpng", fp1, "-O"]
    rawSystem "dot" ["-Tpng", fp2, "-O"]
    rawSystem "explorer" ["."]
  where
    fp1 = "./initial.txt"
    fp2 = "./deleted.txt"
    dotScript1 = dotTree initial
    dotScript2 = dotTree deleted
    initial = fromList ([11,10..0] ++ [12,13] ++ [(-1),(-2)])
    deleted = delMax initial

printTree t = do
    writeFile "./print.txt" (dotTree t)
    rawSystem "dot" ["-Tpng", "./print.txt", "-O"]
    rawSystem "explorer" ["."]

{--------------- Quick Check --------------}

checkBalance xs = (isBalance . fromList) xs

checkDeleteMinBalance xs = (isBalance . deleteMin . fromList) xs

checkDeleteMaxBalance xs = (isBalance . deleteMax . fromList) xs

checkDeleteArbBalance xs
    | null xs = True
    | otherwise = isBalance (delete t k)
  where
    t = fromList xs
    k = head xs
checkOrdered xs = (isOrdered . fromList) xs

checkDeleteMinOrdered xs = (isOrdered . deleteMin . fromList) xs

checkDeleteMaxOrdered xs = (isOrdered . deleteMax . fromList) xs

checkDeleteMinDeleted xs 
    | null nxs || length nxs == 1 = True
    | otherwise = d == m
  where
    nxs = nub xs
    [d] = nxs \\ ys
    m = L.minimum nxs
    ys = (toList . deleteMin . fromList) nxs

checkDeleteMaxDeleted xs
    | null nxs || length nxs == 1 = True
    | otherwise = d == m
  where
    nxs = nub xs
    [d] = nxs \\ ys
    m = L.maximum nxs
    ys = (toList . deleteMax . fromList) nxs

checkAll xs = and [checkBalance xs, checkDeleteMinBalance xs, checkDeleteMaxBalance xs, checkDeleteArbBalance xs, checkOrdered xs, checkDeleteMinOrdered xs, checkDeleteMaxOrdered xs, checkDeleteMinDeleted xs, checkDeleteMaxDeleted xs]

isBalance t = isBlackSame t && isRedSeparate t

isBlackSame t = all (n==) ns
  where
    n:ns = blacks t

blacks = blacks' 0
  where
    blacks' n Leaf = [n+1]
    blacks' n (Node Red l _ r) = blacks' n l ++ blacks' n r
    blacks' n (Node Black l _ r) = blacks' (n+1) l ++ blacks' (n+1) r

isRedSeparate = reds Black

reds _ Leaf = True
reds Red (Node Red _ _ _) = False
reds _ (Node c l _ r) = reds c l && reds c r

isOrdered t = ordered $ toList t
  where
    ordered [] = True
    ordered [_] = True
    ordered (x:y:xys) = x < y && ordered (y:xys)

len xs = length nxs == (length . toList . fromList) nxs where nxs = nub xs