Binary search tree

A BST is a data structure that maintains a set of key-value pairs where the keys have a well-defined order relative to each other. A BST consists of zero or more nodes where each node has 0 to 2 children.

Insertion and lookup are trivial: traverse the tree until you find the node with a matching key or until you hit a dead-end. If the key is smaller than the key of the current node, go down the left child. Otherwise, go down the right child.

Deletion follows a similar process to find the node to be deleted. However, before deleting the node it is replaced with its in-order successor, i.e. the descendant node with the key that is closest yet larger than the sought key. To find this node, go down the right child, then go down the left children there are none left. If no successor could be found, i.e. there is no right child, the left node is moved in stead of the node to be deleted.



module BinTree where

data BinTree k v = BinTree { key   :: k
                           , value :: v
                           , left  :: Maybe (BinTree k v)
                           , right :: Maybe (BinTree k v) }
                           deriving (Show)

insert :: (Ord k) => k -> v -> Maybe (BinTree k v) -> Maybe (BinTree k v)
insert k' v' Nothing = Just $ BinTree k' v' Nothing Nothing
insert k' v' (Just (BinTree k v l r))
  | k' <  k = insert k' v' l >>= \e -> Just $ BinTree k v (Just e) r
  | k' >  k = insert k' v' r >>= \e -> Just $ BinTree k v l (Just e)
  | k' == k = Nothing

retrieve :: (Ord k) => k -> Maybe (BinTree k v) -> Maybe v
retrieve k' Nothing = Nothing
retrieve k' (Just (BinTree k v l r))
  | k' <  k = retrieve k' l
  | k' >  k = retrieve k' r
  | k' == k = Just v

delete :: (Ord k) => k -> Maybe (BinTree k v) -> Maybe (Maybe (BinTree k v))
delete k' Nothing = Nothing
delete k' (Just (BinTree k v l r))
  | k' <  k = delete k' l >>= \e -> Just $ Just $ BinTree k v e r
  | k' >  k = delete k' r >>= \e -> Just $ Just $ BinTree k v l e
  | k' == k = Just $ maybe l
                           (\(k'', v'', t) -> Just $ BinTree k'' v'' l t)
                           (takeFromLeft r)

takeFromLeft :: Maybe (BinTree k v) -> Maybe (k, v, Maybe (BinTree k v))
takeFromLeft Nothing = Nothing
takeFromLeft (Just (BinTree k v Nothing r)) = Just (k, v, r)
takeFromLeft (Just (BinTree k v l r)) =
  takeFromLeft l
  >>= \(k', v', l') -> Just (k', v', Just $ BinTree k v l' r)

prettyPrint :: (Show k, Show v) => Maybe (BinTree k v) -> IO ()
prettyPrint t = prettyPrint' t 0

prettyPrint' :: (Show k, Show v) => Maybe (BinTree k v) -> Int -> IO ()
prettyPrint' Nothing i = return $ ()
prettyPrint' (Just (BinTree k v l r)) i = 
  prettyPrint' r (i + 1)
  >> prettyPrint'' k v i
  >> prettyPrint' l (i + 1)

prettyPrint'' :: (Show k, Show v) => k -> v -> Int -> IO ()
prettyPrint'' k v i = 
  putStr (take i $ repeat ' ')
  >> putStr (show k)
  >> putChar ':'
  >> putStrLn (show v)