Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import Data.Foldable
- import qualified Data.List as L
- data Tree a = Node Integer (Tree a) (Tree a) a | Leaf deriving (Show, Eq)
- data Path = LeftP | RightP
- instance Foldable Tree where
- foldMap f Leaf = mempty
- foldMap f (Node _ l r v) = foldMap f l `mappend` f v `mappend` foldMap f r
- leftRotate :: Tree a -> Tree a
- leftRotate (Node tlvl tleft (Node rlvl rleft rright rdata) tdata) =
- Node rlvl (Node tlvl tleft rleft tdata) rright rdata
- leftRotate tree = tree
- rightRotate (Node tlvl (Node llvl lleft lright ldata) tright tdata) =
- Node llvl lleft (Node tlvl lright tright tdata) ldata
- rightRotate tree = tree
- left :: Tree a -> Tree a
- left Leaf = Leaf
- left (Node _ left _ _) = left
- right :: Tree a -> Tree a
- right Leaf = Leaf
- right (Node _ _ right _) = right
- level :: Tree a -> Integer
- level Leaf = 0
- level (Node lvl _ _ _) = lvl
- isLeaf :: Tree a -> Bool
- isLeaf Leaf = True
- isLeaf _ = False
- mapLevel :: (Integer -> Integer) -> Tree a -> Tree a
- mapLevel _ Leaf = Leaf
- mapLevel f (Node lvl l r d) = Node (f lvl) l r d
- incLevel :: Tree a -> Tree a
- incLevel = mapLevel (1+)
- decLevel :: Tree a -> Tree a
- decLevel = mapLevel (1-)
- skew :: Tree a -> Tree a
- skew tree
- | level tree == (level . left) tree = rightRotate tree
- | otherwise = tree
- split :: Tree a -> Tree a
- split tree
- | level tree == (level . right . right) tree = incLevel . leftRotate $ tree
- | otherwise = tree
- mapAt :: (Tree a -> Tree a) -> [Path] -> Tree a -> Tree a
- mapAt f _ Leaf = f Leaf
- mapAt f [] t = f t
- mapAt f (RightP:xs) (Node lvl l r k) = Node lvl l (mapAt f xs r) k
- mapAt f (LeftP:xs) (Node lvl l r k) = Node lvl (mapAt f xs l) r k
- mapLeft :: (Tree a -> Tree a) -> Tree a -> Tree a
- mapLeft f = mapAt f [LeftP]
- mapRight :: (Tree a -> Tree a) -> Tree a -> Tree a
- mapRight f = mapAt f [RightP]
- insert :: Ord a => a -> Tree a -> Tree a
- insert d tree@(Node lvl l r k)
- | d < k = split . skew . mapLeft (insert d) $ tree
- | d > k = split . skew . mapRight (insert d) $ tree
- | otherwise = tree
- insert d Leaf = Node 1 Leaf Leaf d
- predecessor :: Tree a -> Tree a
- predecessor = goRight . left
- where goRight t
- | isLeaf t = t
- | (isLeaf . left) t && (isLeaf . right) t = t
- | otherwise = goRight . right $ t
- successor :: Tree a -> Tree a
- successor = goLeft . right
- where goLeft t
- | isLeaf t = t
- | (isLeaf . left) t && (isLeaf . right) t = t
- | otherwise = goLeft . left $ t
- deleteRebalance :: Tree a -> Tree a
- deleteRebalance Leaf = Leaf
- deleteRebalance tree@(Node lvl l r k)
- -- If we have some links that jump more than one level.
- | level l < newLevel || level r < newLevel =
- splitRebalance . skewRebalance . mapLevel (return newLevel) . mapRight adjustRightLevel $ tree
- | otherwise = tree
- where
- newLevel = min ((level . left) tree) ((level . right) tree) + 1
- adjustRightLevel = mapLevel (min newLevel)
- splitRebalance tree = foldl (flip $ mapAt split) tree $ reverse (L.tails [RightP])
- skewRebalance tree = foldl (flip $ mapAt skew) tree $ reverse (L.tails [RightP, RightP])
- delete :: Ord a => a -> Tree a -> Tree a
- delete d Leaf = Leaf
- delete d tree@(Node lvl left right v)
- | level tree == 1 && v == d = Leaf
- | d < v = deleteRebalance . mapLeft (delete d) $ tree
- | d > v = deleteRebalance . mapRight (delete d) $ tree
- | isLeaf left = let (Node _ _ _ v) = successor tree in
- deleteRebalance . mapRight (delete v) $ Node lvl left right v
- | otherwise = let (Node _ _ _ v) = predecessor tree in
- deleteRebalance . mapLeft (delete v) $ Node lvl left right v
- validTree :: Tree a -> Bool
- validTree Leaf = True
- validTree tree = (level . left) tree < level tree &&
- (level . right) tree <= level tree &&
- (level . right . right) tree < level tree &&
- validTree (left tree) && validTree (right tree)
Add Comment
Please, Sign In to add comment