SplayTree.hs

import Data.Maybe (isJust)

data Tree a = Leaf | Node (Tree a) a (Tree a)

replace :: a -> Tree a -> Tree a
replace a Leaf = Node Leaf a Leaf
replace a (Node l _ r) = Node l a r

rotateR :: Tree a -> Tree a
rotateR (Node (Node x a y) b z) = Node x a (Node y b z)

rotateL :: Tree a -> Tree a
rotateL (Node x a (Node y b z)) = Node (Node x a y) b z

data Context a = Root | L a (Tree a) (Context a) | R a (Tree a) (Context a)
data Zipper a = Zipper (Tree a) (Context a)

update :: (Tree a -> Tree a) -> Zipper a -> Zipper a
update f (Zipper t c) = Zipper (f t) c

downto :: Ord a => a -> Zipper a -> Zipper a
downto _ z@(Zipper Leaf _) = z
downto a z@(Zipper (Node l b r) c) = case compare a b of
  EQ -> z
  LT -> downto a $ Zipper l (L b r c)
  GT -> downto a $ Zipper r (R b l c)

up :: Zipper a -> Zipper a
up z@(Zipper _ Root) = z
up (Zipper l (L a r c)) = Zipper (Node l a r) c
up (Zipper r (R a l c)) = Zipper (Node l a r) c

fromZipper :: Zipper a -> Tree a
fromZipper (Zipper t Root) = t
fromZipper z = fromZipper $ up z

toZipper :: Tree a -> Zipper a
toZipper = flip Zipper Root

empty :: Tree a
empty = Leaf

splay :: Zipper a -> Zipper a
splay z@(Zipper _ Root)              = z
splay z@(Zipper _ (L _ _ Root))      = update rotateR . up $ z
splay z@(Zipper _ (R _ _ Root))      = update rotateL . up $ z
splay z@(Zipper _ (L _ _ (L _ _ _))) = splay . update rotateR . up . update rotateR . up $ z
splay z@(Zipper _ (L _ _ (R _ _ _))) = splay . update rotateL . up . update rotateR . up $ z
splay z@(Zipper _ (R _ _ (L _ _ _))) = splay . update rotateR . up . update rotateL . up $ z
splay z@(Zipper _ (R _ _ (R _ _ _))) = splay . update rotateL . up . update rotateL . up $ z

locate :: Ord a => a -> Tree a -> Maybe (Tree a)
locate a t = case downto a $ toZipper t of
  Zipper Leaf _ -> Nothing
  z             -> Just . fromZipper . splay . update (replace a) $ z

member :: Ord a => a -> Tree a -> Bool
member a = isJust . locate a

insert :: Ord a => a -> Tree a -> Tree a
insert a = fromZipper . update (replace a) . downto a . toZipper

fromList :: Ord a => [a] -> Tree a
fromList = foldr insert empty

toList :: Tree a -> [a]
toList Leaf = []
toList (Node l a r) = toList l ++ [a] ++ toList r

edges :: Tree a -> [(a, a)]
edges Leaf = []
edges (Node Leaf _ Leaf) = []
edges (Node l@(Node _ b _) a Leaf) = (a, b) : edges l
edges (Node Leaf a r@(Node _ b _)) = (a, b) : edges r
edges (Node l@(Node _ b _) a r@(Node _ c _)) = (a, b) : (a, c) : edges l ++ edges r

dot :: Show a => Tree a -> String
dot t = "digraph tree { " ++ concatMap f (edges t) ++ "}"
  where f (a, b) = show a ++ " -> " ++ show b ++ "; "

main :: IO ()
main = do
  let t = fromList [1, 9, 3, 2, 7, 4, 6, 2, 1, 8, 5]
  Just t <- return $ locate 1 t
  Just t <- return $ locate 2 t
  Just t <- return $ locate 3 t
  Just t <- return $ locate 4 t
  Just t <- return $ locate 5 t
  Just t <- return $ locate 6 t
  Just t <- return $ locate 7 t
  Just t <- return $ locate 8 t
  Just t <- return $ locate 9 t
  putStrLn . dot $ t