Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Toy Splay Tree in Haskell
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
@eraserhd

This comment has been minimized.

Copy link

eraserhd commented Jan 22, 2014

From how I understand it, lines 47 and 50 are wrong. They should read:

splay z@(Zipper _ (L _ _ (L _ _ _))) = splay . update rotateR . update rotateR . up . up $ z

and

splay z@(Zipper _ (R _ _ (R _ _ _))) = splay . update rotateL . update rotateL . up . up $ z

I'm not sure what difference it makes, but the claim is that splaying is different from rotate-to-root, which is what yours does (in yours, there's no need to handle things in pairs, for example).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.