AVL-tree

module AVL where

-- map based on balanced AVL-tree
data AVL k v = Leaf | Node k v (AVL k v) (AVL k v) Int deriving (Show)

-- это в кр не надо, просто функция вывода дерева в консоль
pp :: Show k => AVL k v -> IO ()
pp = (mapM_ putStrLn) . treeIndent
  where
    treeIndent Leaf           = ["-- /-"]
    treeIndent (Node k v lb rb _) =
      ["--" ++ (show k)] ++
      map ("  |" ++) ls ++
      ("  `" ++ r) : map ("   " ++) rs
      where
        (r:rs) = treeIndent $ rb
        ls     = treeIndent $ lb

height :: AVL k v -> Int
height Leaf             = 0
height (Node _ _ _ _ h) = h

maxHeight :: AVL k v -> AVL k v -> Int
maxHeight t1 t2 = 1 + max (height t1) (height t2)

balance :: AVL k v -> Int
balance Leaf = 0
balance (Node _ _ l r _) = height l - height r

rotateLeft :: AVL k v -> AVL k v
rotateLeft (Node kx vx l (Node ky vy rl rr _) _) =
    Node ky vy newLeft rr (maxHeight newLeft rr)
    where newLeft = Node kx vx l rl (maxHeight l rl)

rotateRight :: AVL k v -> AVL k v
rotateRight (Node kx vx (Node ky vy ll lr _) r _) =
    Node ky vy ll newRight (maxHeight ll newRight)
    where newRight = Node kx vx lr r (maxHeight lr r)

rebalance :: AVL k v -> AVL k v
rebalance t@(Node k v l r _)
    | (balance t == -2) && (balance r <= 0) = rotateLeft t -- small left rotation
    | (balance t == 2) && (balance l >= 0) = rotateRight t -- small right rotation
    | (balance t == -2) && (balance r == 1) = rotateLeft (Node k v l (rotateRight r) 0) -- big left rotation
    | (balance t == 2) && (balance l == -1) = rotateRight (Node k v (rotateLeft l) r 0) -- big right rotation
    | otherwise = t

insertAVL :: (Ord k) => k -> v -> AVL k v -> AVL k v
insertAVL k v Leaf = Node k v Leaf Leaf 1
insertAVL k v (Node k' v' l r h)
    | k < k'  = let newLeft = insertAVL k v l in rebalance $ Node k' v' newLeft r (maxHeight newLeft r)
    | k > k'  = let newRight = insertAVL k v r in rebalance $ Node k' v' l newRight (maxHeight l newRight)
    | k == k' = Node k' v l r h

-- примеры работы; можно запустить в консоли pp t1 и посмотреть, что действительно сбалансированное дерево
t1 = foldr (\t -> insertAVL (fst t) (snd t)) Leaf (zip [1..10] [1..10])