Skip to content

Instantly share code, notes, and snippets.

@m2ym
Created December 7, 2012 10:34
Show Gist options
  • Star 2 You must be signed in to star a gist
  • Fork 2 You must be signed in to fork a gist
  • Save m2ym/4232390 to your computer and use it in GitHub Desktop.
Save m2ym/4232390 to your computer and use it in GitHub Desktop.
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
Copy link

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