Guest User

Untitled

a guest
Dec 10th, 2018
101
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. {-# OPTIONS_GHC -O2 -funbox-strict-fields #-}
  2.  
  3. data Tree a = Empty | Node Int !(Tree a) !a !(Tree a) deriving (Show,Eq)
  4.  
  5. size Empty = 0
  6. size (Node s _ _ _) = s
  7.  
  8. update Empty = Empty
  9. update (Node _ l x r) = Node (size l + 1 + size r) l x r
  10.  
  11. uNode l x r = update $ Node 0 l x r
  12.  
  13. insert :: (Ord a) => Tree a -> a -> Tree a
  14. insert t y = findVal (go t y) y
  15.     where go Empty y                        = uNode Empty y Empty
  16.           go t@(Node _ l x r) y | y < x     = uNode (go l y) x r
  17.                                 | x < y     = uNode l x (go r y)
  18.                                 | otherwise = t
  19.  
  20. delete :: (Ord a, Bounded a) => Tree a -> a -> Tree a
  21. delete Empty _         = Empty
  22. delete t x | x == y    = go l r
  23.            | otherwise = tt
  24.     where tt@(Node _ l y r) = findVal t x
  25.           go Empty t        = t
  26.           go s t            = uNode l x t
  27.               where (Node _ l x r) = findVal s maxBound
  28.  
  29. rotate LT (Node _ l x r) (Node _ u y v) = uNode u y (uNode v x r)
  30. rotate GT (Node _ l x r) (Node _ u y v) = uNode (uNode l x u) y v
  31.  
  32. splay :: (Ord a) => [(Tree a,Ordering)] -> Tree a
  33. splay [(x,_)]                               = x
  34. splay ((r,dr):(q,dq):(p,dp):xs) | dp == dq  = splay $ (rotate dq (rotate dp p q) r, dr) : xs
  35.                                 | otherwise = splay $ (rotate dp p (rotate dq q r), dr) : xs
  36. splay ((y,_):(x,dx):_)                      = rotate dx x y
  37.  
  38. findVal :: (Ord a) => Tree a -> a -> Tree a
  39. findVal Empty _            = Empty
  40. findVal n@(Node _ l x r) y = splay $! reverse $! go n y
  41.     where go Empty _                        = []
  42.           go n@(Node _ l x r) y | y < x     = (n,LT) : go l y
  43.                                 | x < y     = (n,GT) : go r y
  44.                                 | otherwise = [(n,EQ)]
  45.  
  46. findKth :: (Ord a) => Tree a -> Int -> Tree a
  47. findKth Empty _            = Empty
  48. findKth n@(Node _ l x r) k = splay $! reverse $! go n k
  49.     where go Empty _                              = []
  50.           go n@(Node _ l x r) k | size l >= k     = (n,LT) : go l k
  51.                                 | size l + 1 == k = [(n,EQ)]
  52.                                 | otherwise       = (n,GT) : go r (k - 1 - size l)
  53.  
  54. getKth :: (Ord a, Show a) => Tree a -> Int -> (Tree a,String)
  55. getKth t k | k > size t = (t,"invalid")
  56.            | otherwise  = (t,show x)
  57.     where (Node _ _ x _) = findKth t k
  58.  
  59. getCount :: (Ord a, Eq a) => Tree a -> a -> (Tree a,String)
  60. getCount Empty _         = (Empty,"0")
  61. getCount t k | x <  k    = (tt, show $ size l + 1)
  62.              | otherwise = (tt, show $ size l)
  63.     where tt@(Node _ l x _) = findVal t k
  64.  
  65. solve :: Tree Int -> [String] -> IO ()
  66. solve _ [] = return ()
  67. solve t (a:b:xs) = do
  68.   let i = read b :: Int
  69.   return ()
  70.   case a of
  71.     "I" -> solve (insert t i) xs
  72.     "D" -> solve (delete t i) xs
  73.     "K" -> do
  74.            let (tt,s) = getKth t i
  75.            putStrLn s
  76.            solve tt xs
  77.     "C" -> do
  78.            let (tt,s) = getCount t i
  79.            putStrLn s
  80.            solve tt xs
  81.     _ -> solve t xs
  82.  
  83. main :: IO ()
  84. main = do
  85.   (_:xs) <- fmap words $ getContents
  86.   solve Empty xs
Add Comment
Please, Sign In to add comment