Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- {-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
- data Tree a = Empty | Node Int !(Tree a) !a !(Tree a) deriving (Show,Eq)
- size Empty = 0
- size (Node s _ _ _) = s
- update Empty = Empty
- update (Node _ l x r) = Node (size l + 1 + size r) l x r
- uNode l x r = update $ Node 0 l x r
- insert :: (Ord a) => Tree a -> a -> Tree a
- insert t y = findVal (go t y) y
- where go Empty y = uNode Empty y Empty
- go t@(Node _ l x r) y | y < x = uNode (go l y) x r
- | x < y = uNode l x (go r y)
- | otherwise = t
- delete :: (Ord a, Bounded a) => Tree a -> a -> Tree a
- delete Empty _ = Empty
- delete t x | x == y = go l r
- | otherwise = tt
- where tt@(Node _ l y r) = findVal t x
- go Empty t = t
- go s t = uNode l x t
- where (Node _ l x r) = findVal s maxBound
- rotate LT (Node _ l x r) (Node _ u y v) = uNode u y (uNode v x r)
- rotate GT (Node _ l x r) (Node _ u y v) = uNode (uNode l x u) y v
- splay :: (Ord a) => [(Tree a,Ordering)] -> Tree a
- splay [(x,_)] = x
- splay ((r,dr):(q,dq):(p,dp):xs) | dp == dq = splay $ (rotate dq (rotate dp p q) r, dr) : xs
- | otherwise = splay $ (rotate dp p (rotate dq q r), dr) : xs
- splay ((y,_):(x,dx):_) = rotate dx x y
- findVal :: (Ord a) => Tree a -> a -> Tree a
- findVal Empty _ = Empty
- findVal n@(Node _ l x r) y = splay $! reverse $! go n y
- where go Empty _ = []
- go n@(Node _ l x r) y | y < x = (n,LT) : go l y
- | x < y = (n,GT) : go r y
- | otherwise = [(n,EQ)]
- findKth :: (Ord a) => Tree a -> Int -> Tree a
- findKth Empty _ = Empty
- findKth n@(Node _ l x r) k = splay $! reverse $! go n k
- where go Empty _ = []
- go n@(Node _ l x r) k | size l >= k = (n,LT) : go l k
- | size l + 1 == k = [(n,EQ)]
- | otherwise = (n,GT) : go r (k - 1 - size l)
- getKth :: (Ord a, Show a) => Tree a -> Int -> (Tree a,String)
- getKth t k | k > size t = (t,"invalid")
- | otherwise = (t,show x)
- where (Node _ _ x _) = findKth t k
- getCount :: (Ord a, Eq a) => Tree a -> a -> (Tree a,String)
- getCount Empty _ = (Empty,"0")
- getCount t k | x < k = (tt, show $ size l + 1)
- | otherwise = (tt, show $ size l)
- where tt@(Node _ l x _) = findVal t k
- solve :: Tree Int -> [String] -> IO ()
- solve _ [] = return ()
- solve t (a:b:xs) = do
- let i = read b :: Int
- return ()
- case a of
- "I" -> solve (insert t i) xs
- "D" -> solve (delete t i) xs
- "K" -> do
- let (tt,s) = getKth t i
- putStrLn s
- solve tt xs
- "C" -> do
- let (tt,s) = getCount t i
- putStrLn s
- solve tt xs
- _ -> solve t xs
- main :: IO ()
- main = do
- (_:xs) <- fmap words $ getContents
- solve Empty xs
Add Comment
Please, Sign In to add comment