Instantly share code, notes, and snippets.

# m2ym/SplayTree.hs Created Dec 7, 2012

 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 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).