gistfile1.hs

-- Description: A red-black tree implementation in Haskell
-- Author: Cedric Van Goethem
-- Version 1, 2013

type Tree a = Node a
data Color = Red | Black deriving (Eq, Show)
data Node a = Node a Color (Node a) (Node a) | Nil deriving (Eq, Show)

tree :: (Ord a) => a -> Tree a
tree x = Node x Black Nil Nil

addValue :: Ord a => a -> Tree a -> Tree a
addValue x = addNode (Node x Red Nil Nil)

addNode :: Ord a => Node a -> Tree a -> Tree a
addNode n t = go $ addNode' n t
  where go (Node v _ l r) = Node v Black l r

addNode' :: Ord a => Node a -> Tree a -> Tree a
addNode' Nil			t 			= t
addNode' n@(Node v1 c1 l1 r1) 	t@(Node v2 c2 l2 r2)
	| v1 < v2 					= balance $ Node v2 c2 (addChild n l2) r2
	| v1 > v2 					= balance $ Node v2 c2 l2 (addChild n r2)
	| otherwise 					= t
	
addChild :: Ord a => Node a -> Node a -> Node a
addChild n Nil 	= n
addChild n t 	= addNode' n t

{- Rotations -}
balance :: Node a -> Node a
balance (Node z Black (Node y Red (Node x Red a b) c) d) = Node y Red (Node x Black a b) (Node z Black c d)
balance (Node z Black (Node x Red a (Node y Red b c)) d) = Node y Red (Node x Black a b) (Node z Black c d)
balance (Node x Black a (Node z Red (Node y Red b c) d)) = Node y Red (Node x Black a b) (Node z Black c d)
balance (Node x Black a (Node y Red b (Node z Red c d))) = Node y Red (Node x Black a b) (Node z Black c d)
balance x = x

{- Getters and setters -}
setColor :: Node a -> Color -> Node a
setColor Nil _ = Nil
setColor (Node v _ x y) c = Node v c x y

getLeft :: Node a -> Node a
getLeft (Node _ _ x _) = x

getRight :: Node a -> Node a
getRight (Node _ _ _ y) = y

setLeft :: Node a -> Node a -> Node a
setLeft (Node v c x y) o = Node v c o y

setRight :: Node a -> Node a -> Node a
setRight (Node v c x y) o = Node v c x o