Skip to content

Instantly share code, notes, and snippets.

@hardentoo
Forked from m2ym/SplayTree.hs
Created May 26, 2018 01:05
Show Gist options
  • Save hardentoo/a3349dfa6d4d8e093a07bcf5d8f8dad9 to your computer and use it in GitHub Desktop.
Save hardentoo/a3349dfa6d4d8e093a07bcf5d8f8dad9 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment