{-
     Red black trees
     Author:  Robin Cockett
     Date: 22 Sept 2005

     The invariance under insertion is that
     (1) no red node has a red child
     (2) on any path from root to tip has the same number of black
         nodes.
 -}

module RedBlack (insert,inserts,member)
where
     
 data Colour = Red|Black
     deriving (Show, Eq, Enum, Ord)

 data RBTree a = RBTip | RBNode Colour (RBTree a) a (RBTree a)
     deriving (Show, Eq)

 insert:: Ord a => a -> (RBTree a) -> (RBTree a)
 insert x RBTip = RBNode Black RBTip x RBTip
 insert x t = maketopblack (ins x t) where
      
      maketopblack (RBNode Red t1 x t2) = (RBNode Black t1 x t2)
      maketopblack t = t

      ins x RBTip = RBNode Red RBTip x RBTip
      ins x (RBNode colour t1 y t2)
          | x < y = balance colour (ins x t1) y t2
          | x > y = balance colour t1 y (ins x t2)
          | otherwise = RBNode colour t1 y t2

      balance Black (RBNode Red (RBNode Red t11 x1 t12) y t21) x2 t22 =
               RBNode Red (RBNode Black t11 x1 t12) y (RBNode Black t21 x2 t22)
      balance Black (RBNode Red t11 x1 (RBNode Red t12 y t21)) x2 t22 =
               RBNode Red (RBNode Black t11 x1 t12) y (RBNode Black t21 x2 t22)
      balance Black t11 x1 (RBNode Red (RBNode Red t12 y t21) x2 t22) =
               RBNode Red (RBNode Black t11 x1 t12) y (RBNode Black t21 x2 t22)
      balance Black t11 x1 (RBNode Red t12 y (RBNode Red t21 x2 t22)) =
               RBNode Red (RBNode Black t11 x1 t12) y (RBNode Black t21 x2 t22)
      balance colour t1 x1 t2 = RBNode colour t1 x1 t2

 inserts:: Ord a => [a] -> (RBTree a) -> (RBTree a)
 inserts [] t = t
 inserts (a:as) t = insert a (inserts as t)

 member:: Ord a => a -> (RBTree a) -> Bool
 member a RBTip = False
 member a (RBNode c t1 b t2)
       | a < b = member a t1
       | a > b = member a t2
       | otherwise = True